okbot_agent/
chain.rs

1//! Ordered fallback across backends, with tool-policy enforcement and a concurrency limit.
2
3use std::sync::Arc;
4
5use okbot_core::{AgentBackend, AgentCaps, AgentError, AgentReply, AgentTurn, async_trait};
6use tokio::sync::Semaphore;
7
8/// Tries backends in order until one answers.
9///
10/// - A backend that cannot enforce the turn's tool policy is skipped, unless the chain was built
11///   with [`AgentChain::allow_unenforced_tools`] (development only: the agent may then use every
12///   tool its CLI has, including a shell).
13/// - Errors that [`AgentError::should_fall_back`] move on to the next backend; others stop.
14/// - At most `max_concurrent` turns run at once across the whole chain.
15pub struct AgentChain {
16    backends: Vec<Arc<dyn AgentBackend>>,
17    name: String,
18    permits: Arc<Semaphore>,
19    machine: Option<crate::MachineSlots>,
20    allow_unenforced: bool,
21}
22
23impl std::fmt::Debug for AgentChain {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("AgentChain")
26            .field("name", &self.name)
27            .field("allow_unenforced", &self.allow_unenforced)
28            .finish_non_exhaustive()
29    }
30}
31
32impl AgentChain {
33    /// A chain over `backends` (tried in order) running at most `max_concurrent` turns at once.
34    pub fn new(backends: Vec<Arc<dyn AgentBackend>>, max_concurrent: usize) -> Self {
35        let names: Vec<&str> = backends.iter().map(|b| b.name()).collect();
36        Self {
37            name: names.join(","),
38            backends,
39            permits: Arc::new(Semaphore::new(max_concurrent.max(1))),
40            machine: None,
41            allow_unenforced: false,
42        }
43    }
44
45    /// Lets backends run turns whose tool policy they cannot enforce. For local development with
46    /// a trusted user only; never for a bot reachable by others.
47    pub fn allow_unenforced_tools(mut self, allow: bool) -> Self {
48        if allow {
49            tracing::warn!(
50                chain = %self.name,
51                "tool policies are NOT enforced: the agent may run shell commands (development only)"
52            );
53        }
54        self.allow_unenforced = allow;
55        self
56    }
57
58    /// Also waits for a machine-wide slot before each turn (shared with other bots).
59    pub fn machine_slots(mut self, slots: crate::MachineSlots) -> Self {
60        self.machine = Some(slots);
61        self
62    }
63
64    /// The backends, in order.
65    pub fn backends(&self) -> &[Arc<dyn AgentBackend>] {
66        &self.backends
67    }
68}
69
70#[async_trait]
71impl AgentBackend for AgentChain {
72    fn name(&self) -> &str {
73        &self.name
74    }
75
76    /// Native features shared by every backend; any backend's enforceable tool policies.
77    fn caps(&self) -> AgentCaps {
78        let all: Vec<AgentCaps> = self.backends.iter().map(|b| b.caps()).collect();
79        let every = |f: fn(&AgentCaps) -> bool| !all.is_empty() && all.iter().all(f);
80        let mut tool_policies = Vec::new();
81        for caps in &all {
82            for policy in &caps.tool_policies {
83                if !tool_policies.contains(policy) {
84                    tool_policies.push(*policy);
85                }
86            }
87        }
88        AgentCaps {
89            system_prompt: every(|c| c.system_prompt),
90            mcp: every(|c| c.mcp),
91            tool_policies,
92            resume: every(|c| c.resume),
93            stream: every(|c| c.stream),
94        }
95    }
96
97    async fn run(&self, turn: AgentTurn) -> Result<AgentReply, AgentError> {
98        let _permit = self
99            .permits
100            .acquire()
101            .await
102            .map_err(|_| AgentError::Unavailable("agent chain is shutting down".into()))?;
103        let _slot =
104            match &self.machine {
105                Some(slots) => Some(slots.acquire().await.map_err(|e| {
106                    AgentError::Unavailable(format!("machine-wide agent limit: {e}"))
107                })?),
108                None => None,
109            };
110        let mut last_error = AgentError::Unavailable("no agent backend configured".into());
111        for backend in &self.backends {
112            if !self.allow_unenforced && !backend.caps().enforces(turn.tools) {
113                tracing::debug!(backend = backend.name(), policy = ?turn.tools, "skipped: cannot enforce tool policy");
114                last_error = AgentError::Unsupported(format!(
115                    "{} cannot enforce tool policy {:?}",
116                    backend.name(),
117                    turn.tools
118                ));
119                continue;
120            }
121            match backend.run(turn.clone()).await {
122                Ok(reply) => return Ok(reply),
123                Err(error) if error.should_fall_back() => {
124                    tracing::warn!(backend = backend.name(), %error, "agent failed, trying the next one");
125                    last_error = error;
126                }
127                Err(error) => return Err(error),
128            }
129        }
130        Err(last_error)
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use okbot_core::ToolPolicy;
138    use okbot_testkit::{FakeAgent, turn};
139
140    fn chain(backends: Vec<Arc<dyn AgentBackend>>) -> AgentChain {
141        AgentChain::new(backends, 2)
142    }
143
144    #[tokio::test]
145    async fn falls_back_on_retryable_errors() {
146        let a = Arc::new(FakeAgent::failing("a", true).enforcing(&[ToolPolicy::Web]));
147        let b = Arc::new(FakeAgent::echo("b").enforcing(&[ToolPolicy::Web]));
148        let reply = chain(vec![a.clone(), b.clone()])
149            .run(turn("hi"))
150            .await
151            .unwrap();
152        assert_eq!(reply.backend, "b");
153        assert_eq!((a.calls(), b.calls()), (1, 1));
154    }
155
156    #[tokio::test]
157    async fn stops_on_permanent_errors() {
158        let a = Arc::new(FakeAgent::failing("a", false).enforcing(&[ToolPolicy::Web]));
159        let b = Arc::new(FakeAgent::echo("b").enforcing(&[ToolPolicy::Web]));
160        assert!(chain(vec![a, b.clone()]).run(turn("hi")).await.is_err());
161        assert_eq!(b.calls(), 0);
162    }
163
164    #[tokio::test]
165    async fn skips_backends_that_cannot_enforce_the_policy() {
166        let a = Arc::new(FakeAgent::echo("a"));
167        let b = Arc::new(FakeAgent::echo("b").enforcing(&[ToolPolicy::Web]));
168        let reply = chain(vec![a.clone(), b]).run(turn("hi")).await.unwrap();
169        assert_eq!(reply.backend, "b");
170        assert_eq!(a.calls(), 0);
171
172        let only_unenforced = chain(vec![Arc::new(FakeAgent::echo("a"))]);
173        assert!(matches!(
174            only_unenforced.run(turn("hi")).await,
175            Err(AgentError::Unsupported(_))
176        ));
177        let dev = chain(vec![Arc::new(FakeAgent::echo("a"))]).allow_unenforced_tools(true);
178        assert_eq!(dev.run(turn("hi")).await.unwrap().backend, "a");
179    }
180
181    #[tokio::test]
182    async fn caps_combine() {
183        let a = Arc::new(FakeAgent::echo("a").enforcing(&[ToolPolicy::Web]));
184        let b = Arc::new(FakeAgent::echo("b").enforcing(&[ToolPolicy::McpOnly]));
185        let caps = chain(vec![a, b]).caps();
186        assert!(caps.enforces(ToolPolicy::Web) && caps.enforces(ToolPolicy::McpOnly));
187        assert!(!caps.system_prompt);
188    }
189}
okbot_agent/chain.rs

okbot_agent/
chain.rs
