Skip to main content

monoloop_loop/transaction/mcp/
binding.rs

1//! Capability tokens and transaction MCP route table.
2
3use super::handler::TransactionMcpHandler;
4use crate::transaction::dispatcher::TransactionToolDispatcher;
5use crate::transaction::resolved_tools::ResolvedToolSet;
6use monoloop_connector::McpServerDescriptor;
7use monoloop_contracts::{ExchangeId, TransactionId};
8use rand::TryRngCore;
9use std::collections::HashMap;
10use std::fmt;
11use std::sync::atomic::{AtomicU8, Ordering};
12use std::sync::{Arc, Mutex};
13
14const STATE_PENDING: u8 = 0;
15const STATE_ACTIVE: u8 = 1;
16const STATE_REVOKED: u8 = 2;
17
18/// 256-bit unguessable capability token (hex in URLs; redacted in diagnostics).
19#[derive(Clone, PartialEq, Eq, Hash)]
20pub struct CapabilityToken {
21    bytes: [u8; 32],
22}
23
24impl CapabilityToken {
25    /// Generate via OS CSPRNG.
26    pub fn generate() -> Result<Self, McpInstallError> {
27        let mut bytes = [0u8; 32];
28        rand::rngs::OsRng
29            .try_fill_bytes(&mut bytes)
30            .map_err(|_| McpInstallError::EntropyUnavailable)?;
31        if bytes == [0u8; 32] {
32            return Err(McpInstallError::EntropyUnavailable);
33        }
34        Ok(Self { bytes })
35    }
36
37    /// Parse a 64-char lowercase hex token (URL segment).
38    pub fn from_hex(s: &str) -> Option<Self> {
39        if s.len() != 64 || !s.bytes().all(|b| b.is_ascii_hexdigit()) {
40            return None;
41        }
42        let mut bytes = [0u8; 32];
43        let chars = s.as_bytes();
44        for (i, slot) in bytes.iter_mut().enumerate() {
45            let hi = hex_nibble(chars[i * 2])?;
46            let lo = hex_nibble(chars[i * 2 + 1])?;
47            *slot = (hi << 4) | lo;
48        }
49        Some(Self { bytes })
50    }
51
52    /// Lowercase hex (64 chars) for URL path segments only.
53    pub fn to_hex(&self) -> String {
54        let mut out = String::with_capacity(64);
55        for b in &self.bytes {
56            out.push(hex_char(b >> 4));
57            out.push(hex_char(b & 0x0f));
58        }
59        out
60    }
61}
62
63impl fmt::Debug for CapabilityToken {
64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65        f.write_str("CapabilityToken(<redacted>)")
66    }
67}
68
69fn hex_nibble(b: u8) -> Option<u8> {
70    match b {
71        b'0'..=b'9' => Some(b - b'0'),
72        b'a'..=b'f' => Some(b - b'a' + 10),
73        b'A'..=b'F' => Some(b - b'A' + 10),
74        _ => None,
75    }
76}
77
78fn hex_char(n: u8) -> char {
79    match n {
80        0..=9 => (b'0' + n) as char,
81        10..=15 => (b'a' + n - 10) as char,
82        _ => '0',
83    }
84}
85
86/// Public lifecycle state of a capability route.
87#[derive(Clone, Copy, Debug, PartialEq, Eq)]
88pub enum McpBindingState {
89    /// Installed but not yet activated (SessionKey not claimed / refresh incomplete).
90    Pending,
91    /// Ready for tools/list and tools/call.
92    Active,
93    /// Revoked; route removed or rejects all traffic.
94    Revoked,
95}
96
97/// Errors installing or mutating MCP routes.
98#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
99pub enum McpInstallError {
100    /// OS CSPRNG failed.
101    #[error("capability entropy unavailable")]
102    EntropyUnavailable,
103    /// Unknown token.
104    #[error("unknown MCP capability")]
105    UnknownCapability,
106    /// Route not in expected state.
107    #[error("MCP capability state conflict")]
108    StateConflict,
109    /// Route table capacity exceeded.
110    #[error("MCP route table full")]
111    CapacityExceeded,
112    /// Invalid descriptor construction.
113    #[error("invalid MCP descriptor")]
114    InvalidDescriptor,
115}
116
117/// One transaction MCP binding (pending or active).
118pub struct McpBinding {
119    /// Capability token.
120    pub token: CapabilityToken,
121    /// Owning transaction.
122    pub transaction_id: TransactionId,
123    /// Shared state flag.
124    state: Arc<AtomicU8>,
125    /// Handler for this binding.
126    pub handler: TransactionMcpHandler,
127    /// Dispatcher (same instance as model path when both used).
128    pub dispatcher: Arc<TransactionToolDispatcher>,
129    /// Resolved tools projection.
130    pub tools: ResolvedToolSet,
131}
132
133impl McpBinding {
134    /// Current lifecycle state.
135    pub fn state(&self) -> McpBindingState {
136        match self.state.load(Ordering::SeqCst) {
137            STATE_PENDING => McpBindingState::Pending,
138            STATE_ACTIVE => McpBindingState::Active,
139            _ => McpBindingState::Revoked,
140        }
141    }
142
143    /// Whether tools/list and tools/call are allowed.
144    pub fn is_active(&self) -> bool {
145        self.state.load(Ordering::SeqCst) == STATE_ACTIVE
146    }
147}
148
149impl fmt::Debug for McpBinding {
150    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
151        f.debug_struct("McpBinding")
152            .field("token", &self.token)
153            .field("transaction_id", &self.transaction_id)
154            .field("state", &self.state())
155            .field("tool_count", &self.tools.len())
156            .finish()
157    }
158}
159
160/// Handle returned when a pending binding is created.
161#[derive(Clone)]
162pub struct PendingMcpBinding {
163    /// Capability token (for activation/revoke; never log).
164    pub token: CapabilityToken,
165    /// Redacted descriptor for SessionAdapter install.
166    pub descriptor: McpServerDescriptor,
167    /// Transaction id.
168    pub transaction_id: TransactionId,
169    /// Dispatcher to rebind after authoritative session claim (D-026).
170    pub dispatcher: Arc<TransactionToolDispatcher>,
171}
172
173impl fmt::Debug for PendingMcpBinding {
174    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175        f.debug_struct("PendingMcpBinding")
176            .field("token", &self.token)
177            .field("descriptor", &self.descriptor)
178            .field("transaction_id", &self.transaction_id)
179            .field("dispatcher_session", &self.dispatcher.session_key())
180            .finish()
181    }
182}
183
184/// Bounded in-memory capability route table.
185pub struct McpRouteTable {
186    max_routes: usize,
187    routes: Mutex<HashMap<CapabilityToken, Arc<McpBinding>>>,
188}
189
190impl McpRouteTable {
191    /// Create with a maximum concurrent route count.
192    pub fn new(max_routes: usize) -> Arc<Self> {
193        Arc::new(Self {
194            max_routes: max_routes.max(1),
195            routes: Mutex::new(HashMap::new()),
196        })
197    }
198
199    /// Number of live (pending or active) routes.
200    pub fn len(&self) -> usize {
201        self.routes.lock().map(|m| m.len()).unwrap_or(0)
202    }
203
204    /// Whether empty.
205    pub fn is_empty(&self) -> bool {
206        self.len() == 0
207    }
208
209    /// Install a pending binding and return the redacted descriptor URL.
210    pub fn install_pending(
211        self: &Arc<Self>,
212        transaction_id: TransactionId,
213        tools: ResolvedToolSet,
214        dispatcher: Arc<TransactionToolDispatcher>,
215        exchange_id: ExchangeId,
216        base_url: &str,
217    ) -> Result<PendingMcpBinding, McpInstallError> {
218        self.install_pending_with_deadline(
219            transaction_id,
220            tools,
221            dispatcher,
222            exchange_id,
223            base_url,
224            std::time::Instant::now() + std::time::Duration::from_secs(365 * 24 * 3600),
225        )
226    }
227
228    /// Install with the live transaction absolute Instant (caps MCP tool budgets).
229    pub fn install_pending_with_deadline(
230        self: &Arc<Self>,
231        transaction_id: TransactionId,
232        tools: ResolvedToolSet,
233        dispatcher: Arc<TransactionToolDispatcher>,
234        exchange_id: ExchangeId,
235        base_url: &str,
236        transaction_deadline: std::time::Instant,
237    ) -> Result<PendingMcpBinding, McpInstallError> {
238        let token = CapabilityToken::generate()?;
239        let state = Arc::new(AtomicU8::new(STATE_PENDING));
240        let handler = TransactionMcpHandler::new(
241            Arc::clone(&state),
242            tools.clone(),
243            Arc::clone(&dispatcher),
244            transaction_id,
245            exchange_id,
246            transaction_deadline,
247        );
248        let binding = Arc::new(McpBinding {
249            token: token.clone(),
250            transaction_id,
251            state,
252            handler,
253            dispatcher: Arc::clone(&dispatcher),
254            tools,
255        });
256
257        {
258            let mut map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
259            if map.len() >= self.max_routes {
260                return Err(McpInstallError::CapacityExceeded);
261            }
262            map.insert(token.clone(), binding);
263        }
264
265        let url = format!("{}/mcp/{}", base_url.trim_end_matches('/'), token.to_hex());
266        let descriptor = McpServerDescriptor::try_new("monoloop", "2024-11-05", url)
267            .map_err(|_| McpInstallError::InvalidDescriptor)?;
268
269        Ok(PendingMcpBinding {
270            token,
271            descriptor,
272            transaction_id,
273            dispatcher,
274        })
275    }
276
277    /// Activate a pending route after SessionKey claim / MCP refresh.
278    pub fn activate(&self, token: &CapabilityToken) -> Result<(), McpInstallError> {
279        let map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
280        let binding = map.get(token).ok_or(McpInstallError::UnknownCapability)?;
281        let prev = binding.state.compare_exchange(
282            STATE_PENDING,
283            STATE_ACTIVE,
284            Ordering::SeqCst,
285            Ordering::SeqCst,
286        );
287        match prev {
288            Ok(_) => Ok(()),
289            Err(STATE_ACTIVE) => Ok(()), // idempotent activate
290            Err(_) => Err(McpInstallError::StateConflict),
291        }
292    }
293
294    /// Revoke and remove a route. Idempotent for unknown tokens.
295    pub fn revoke(&self, token: &CapabilityToken) -> bool {
296        let mut map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
297        if let Some(binding) = map.remove(token) {
298            binding.state.store(STATE_REVOKED, Ordering::SeqCst);
299            true
300        } else {
301            false
302        }
303    }
304
305    /// Revoke every route (shutdown). Returns hex tokens that were live.
306    pub fn revoke_all(&self) -> Vec<String> {
307        let mut map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
308        let mut tokens = Vec::with_capacity(map.len());
309        for (token, binding) in map.drain() {
310            binding.state.store(STATE_REVOKED, Ordering::SeqCst);
311            tokens.push(token.to_hex());
312        }
313        tokens
314    }
315
316    /// Lookup live binding by token hex (from URL).
317    pub fn get_by_hex(&self, hex: &str) -> Option<Arc<McpBinding>> {
318        let token = CapabilityToken::from_hex(hex)?;
319        self.get(&token)
320    }
321
322    /// Lookup by token.
323    pub fn get(&self, token: &CapabilityToken) -> Option<Arc<McpBinding>> {
324        self.routes
325            .lock()
326            .unwrap_or_else(|e| e.into_inner())
327            .get(token)
328            .cloned()
329    }
330
331    /// State for a token if present.
332    pub fn state_of(&self, token: &CapabilityToken) -> Option<McpBindingState> {
333        self.get(token).map(|b| b.state())
334    }
335}