Skip to main content

bitrouter_core/routers/
admin.rs

1//! Admin types and traits for runtime management of routes, tools, and agents.
2//!
3//! Provides extension traits that layer admin (mutation / inspection) capabilities
4//! on top of the read-only discovery traits in [`registry`](super::registry):
5//!
6//! | Entity  | Discovery trait    | Admin trait            |
7//! |---------|--------------------|------------------------|
8//! | Models  | `RoutingTable`     | `AdminRoutingTable`    |
9//! | Tools   | `ToolRegistry`     | `AdminToolRegistry`    |
10//! | Agents  | `AgentRegistry`    | `AdminAgentRegistry`   |
11
12use std::collections::HashMap;
13use std::future::Future;
14
15use serde::{Deserialize, Serialize};
16
17use crate::errors::Result;
18
19use super::registry::{AgentRegistry, ToolRegistry};
20use super::routing_table::RoutingTable;
21
22/// A single endpoint in a dynamic route.
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct RouteEndpoint {
25    /// Provider name (must exist in the providers section or built-ins).
26    pub provider: String,
27    /// The upstream model ID to send to this provider.
28    pub model_id: String,
29}
30
31/// Strategy for distributing requests across multiple endpoints.
32#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
33#[serde(rename_all = "snake_case")]
34pub enum RouteStrategy {
35    /// Try endpoints in declared order.
36    #[default]
37    Priority,
38    /// Distribute requests evenly via round-robin.
39    LoadBalance,
40}
41
42/// A dynamically-configured route definition.
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct DynamicRoute {
45    /// The virtual model name (e.g. "research", "fast").
46    pub model: String,
47    /// Routing strategy across endpoints.
48    #[serde(default)]
49    pub strategy: RouteStrategy,
50    /// One or more upstream endpoints to route to.
51    pub endpoints: Vec<RouteEndpoint>,
52}
53
54/// Extension trait for routing tables that support runtime route management.
55///
56/// Implementations store dynamic routes separately from config-defined routes.
57/// Dynamic routes take precedence during resolution.
58pub trait AdminRoutingTable: RoutingTable {
59    /// Create or update a dynamic route.
60    fn add_route(&self, route: DynamicRoute) -> Result<()>;
61
62    /// Remove a dynamically-added route. Returns `true` if the route existed.
63    ///
64    /// Config-defined routes cannot be removed.
65    fn remove_route(&self, model: &str) -> Result<bool>;
66
67    /// List all dynamically-added routes.
68    fn list_dynamic_routes(&self) -> Vec<DynamicRoute>;
69}
70
71// ── Tool admin ──────────────────────────────────────────────────────
72
73/// Allow/deny filter applied to an upstream tool server.
74///
75/// When both `allow` and `deny` are set, deny takes precedence.
76#[derive(Debug, Clone, Default, Serialize, Deserialize)]
77pub struct ToolFilter {
78    /// If set, only tools whose un-namespaced name appears in this list are visible.
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub allow: Option<Vec<String>>,
81    /// Tools whose un-namespaced name appears in this list are hidden.
82    #[serde(default, skip_serializing_if = "Option::is_none")]
83    pub deny: Option<Vec<String>>,
84}
85
86impl ToolFilter {
87    /// Returns `true` if `tool_name` (un-namespaced) passes this filter.
88    pub fn accepts(&self, tool_name: &str) -> bool {
89        if let Some(deny) = &self.deny
90            && deny.iter().any(|d| d == tool_name)
91        {
92            return false;
93        }
94        if let Some(allow) = &self.allow {
95            return allow.iter().any(|a| a == tool_name);
96        }
97        true
98    }
99}
100
101/// Action taken when a parameter violates restrictions.
102#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
103#[serde(rename_all = "snake_case")]
104pub enum ParamViolationAction {
105    /// Remove the parameter silently, proceed with call.
106    Strip,
107    /// Reject the entire tool call.
108    #[default]
109    Reject,
110}
111
112/// Restriction rules for a single tool's parameters.
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct ParamRule {
115    /// Parameters to deny. Deny takes precedence over allow.
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub deny: Option<Vec<String>>,
118    /// If set, only these parameters are allowed through.
119    #[serde(default, skip_serializing_if = "Option::is_none")]
120    pub allow: Option<Vec<String>>,
121    /// What to do when a restricted parameter is found.
122    #[serde(default)]
123    pub action: ParamViolationAction,
124}
125
126/// Per-server parameter restrictions applied before forwarding tool calls.
127#[derive(Debug, Clone, Default, Serialize, Deserialize)]
128pub struct ParamRestrictions {
129    /// Per-tool parameter rules. Keys are un-namespaced tool names.
130    #[serde(default)]
131    pub rules: HashMap<String, ParamRule>,
132}
133
134impl ParamRestrictions {
135    /// Validate and optionally mutate tool call arguments.
136    ///
137    /// Returns `Ok(())` if allowed (possibly with stripped params).
138    /// Returns `Err` if a parameter is denied and action is `Reject`.
139    pub fn check(
140        &self,
141        tool_name: &str,
142        arguments: &mut Option<serde_json::Map<String, serde_json::Value>>,
143    ) -> Result<()> {
144        let Some(rule) = self.rules.get(tool_name) else {
145            return Ok(());
146        };
147        let Some(args) = arguments.as_mut() else {
148            return Ok(());
149        };
150
151        // Deny list takes precedence.
152        if let Some(deny) = &rule.deny {
153            let denied: Vec<String> = args
154                .keys()
155                .filter(|k| deny.iter().any(|d| d == *k))
156                .cloned()
157                .collect();
158            for key in &denied {
159                match rule.action {
160                    ParamViolationAction::Reject => {
161                        return Err(crate::errors::BitrouterError::invalid_request(
162                            None,
163                            format!("parameter '{key}' denied on tool '{tool_name}'"),
164                            None,
165                        ));
166                    }
167                    ParamViolationAction::Strip => {
168                        args.remove(key);
169                    }
170                }
171            }
172        }
173
174        // Allow list: reject/strip any key NOT in the list.
175        if let Some(allow) = &rule.allow {
176            let disallowed: Vec<String> = args
177                .keys()
178                .filter(|k| !allow.iter().any(|a| a == *k))
179                .cloned()
180                .collect();
181            for key in &disallowed {
182                match rule.action {
183                    ParamViolationAction::Reject => {
184                        return Err(crate::errors::BitrouterError::invalid_request(
185                            None,
186                            format!("parameter '{key}' denied on tool '{tool_name}'"),
187                            None,
188                        ));
189                    }
190                    ParamViolationAction::Strip => {
191                        args.remove(key);
192                    }
193                }
194            }
195        }
196
197        Ok(())
198    }
199}
200
201/// Metadata about a connected upstream tool server.
202#[derive(Debug, Clone, Serialize)]
203pub struct ToolUpstreamEntry {
204    /// Server name.
205    pub name: String,
206    /// Number of tools currently visible (after filtering).
207    pub tool_count: usize,
208    /// Active tool filter, if any.
209    #[serde(skip_serializing_if = "Option::is_none")]
210    pub filter: Option<ToolFilter>,
211    /// Active parameter restrictions, if any.
212    #[serde(skip_serializing_if = "Option::is_none")]
213    pub param_restrictions: Option<ParamRestrictions>,
214}
215
216/// Admin interface for managing tool registries at runtime.
217///
218/// Parallel to [`AdminRoutingTable`] for models. Extends [`ToolRegistry`]
219/// with methods for inspecting upstreams and updating filters and parameter
220/// restrictions without requiring config rewrites or daemon restarts.
221pub trait AdminToolRegistry: ToolRegistry {
222    /// List all upstream tool servers with their current state.
223    fn list_upstreams(&self) -> impl Future<Output = Vec<ToolUpstreamEntry>> + Send;
224    /// List all configured access groups.
225    fn list_groups(&self) -> impl Future<Output = HashMap<String, Vec<String>>> + Send;
226    /// Update the tool filter for a specific upstream server.
227    fn update_filter(
228        &self,
229        server: &str,
230        filter: Option<ToolFilter>,
231    ) -> impl Future<Output = Result<()>> + Send;
232    /// Update parameter restrictions for a specific upstream server.
233    fn update_param_restrictions(
234        &self,
235        server: &str,
236        restrictions: ParamRestrictions,
237    ) -> impl Future<Output = Result<()>> + Send;
238}
239
240// ── Agent admin ─────────────────────────────────────────────────────
241
242/// Metadata about a connected upstream agent.
243#[derive(Debug, Clone, Serialize)]
244pub struct AgentUpstreamEntry {
245    /// Agent name.
246    pub name: String,
247    /// Upstream agent URL.
248    pub url: String,
249    /// Whether the upstream connection is active.
250    pub connected: bool,
251}
252
253/// Trait for providing upstream agent metadata.
254///
255/// Implemented by protocol-specific registries (e.g. A2A) that can report
256/// connection-level details not captured in the generic [`AgentEntry`].
257pub trait AgentUpstreamSource: Send + Sync {
258    /// List upstream agents with connection status.
259    fn list_upstreams(&self) -> impl Future<Output = Vec<AgentUpstreamEntry>> + Send;
260}
261
262impl<T: AgentUpstreamSource> AgentUpstreamSource for std::sync::Arc<T> {
263    async fn list_upstreams(&self) -> Vec<AgentUpstreamEntry> {
264        (**self).list_upstreams().await
265    }
266}
267
268/// Admin interface for inspecting agent registries at runtime.
269///
270/// Parallel to [`AdminRoutingTable`] for models. Extends [`AgentRegistry`]
271/// with methods for inspecting upstream agent connections.
272pub trait AdminAgentRegistry: AgentRegistry {
273    /// List all upstream agents with connection status.
274    fn list_upstreams(&self) -> impl Future<Output = Vec<AgentUpstreamEntry>> + Send;
275}