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.
161pub struct PendingMcpBinding {
162    /// Capability token (for activation/revoke; never log).
163    pub token: CapabilityToken,
164    /// Redacted descriptor for SessionAdapter install.
165    pub descriptor: McpServerDescriptor,
166    /// Transaction id.
167    pub transaction_id: TransactionId,
168}
169
170impl fmt::Debug for PendingMcpBinding {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        f.debug_struct("PendingMcpBinding")
173            .field("token", &self.token)
174            .field("descriptor", &self.descriptor)
175            .field("transaction_id", &self.transaction_id)
176            .finish()
177    }
178}
179
180/// Bounded in-memory capability route table.
181pub struct McpRouteTable {
182    max_routes: usize,
183    routes: Mutex<HashMap<CapabilityToken, Arc<McpBinding>>>,
184}
185
186impl McpRouteTable {
187    /// Create with a maximum concurrent route count.
188    pub fn new(max_routes: usize) -> Arc<Self> {
189        Arc::new(Self {
190            max_routes: max_routes.max(1),
191            routes: Mutex::new(HashMap::new()),
192        })
193    }
194
195    /// Number of live (pending or active) routes.
196    pub fn len(&self) -> usize {
197        self.routes.lock().map(|m| m.len()).unwrap_or(0)
198    }
199
200    /// Whether empty.
201    pub fn is_empty(&self) -> bool {
202        self.len() == 0
203    }
204
205    /// Install a pending binding and return the redacted descriptor URL.
206    pub fn install_pending(
207        self: &Arc<Self>,
208        transaction_id: TransactionId,
209        tools: ResolvedToolSet,
210        dispatcher: Arc<TransactionToolDispatcher>,
211        exchange_id: ExchangeId,
212        base_url: &str,
213    ) -> Result<PendingMcpBinding, McpInstallError> {
214        let token = CapabilityToken::generate()?;
215        let state = Arc::new(AtomicU8::new(STATE_PENDING));
216        let handler = TransactionMcpHandler::new(
217            Arc::clone(&state),
218            tools.clone(),
219            Arc::clone(&dispatcher),
220            transaction_id,
221            exchange_id,
222        );
223        let binding = Arc::new(McpBinding {
224            token: token.clone(),
225            transaction_id,
226            state,
227            handler,
228            dispatcher,
229            tools,
230        });
231
232        {
233            let mut map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
234            if map.len() >= self.max_routes {
235                return Err(McpInstallError::CapacityExceeded);
236            }
237            map.insert(token.clone(), binding);
238        }
239
240        let url = format!("{}/mcp/{}", base_url.trim_end_matches('/'), token.to_hex());
241        let descriptor = McpServerDescriptor::try_new("monoloop", "2024-11-05", url)
242            .map_err(|_| McpInstallError::InvalidDescriptor)?;
243
244        Ok(PendingMcpBinding {
245            token,
246            descriptor,
247            transaction_id,
248        })
249    }
250
251    /// Activate a pending route after SessionKey claim / MCP refresh.
252    pub fn activate(&self, token: &CapabilityToken) -> Result<(), McpInstallError> {
253        let map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
254        let binding = map.get(token).ok_or(McpInstallError::UnknownCapability)?;
255        let prev = binding.state.compare_exchange(
256            STATE_PENDING,
257            STATE_ACTIVE,
258            Ordering::SeqCst,
259            Ordering::SeqCst,
260        );
261        match prev {
262            Ok(_) => Ok(()),
263            Err(STATE_ACTIVE) => Ok(()), // idempotent activate
264            Err(_) => Err(McpInstallError::StateConflict),
265        }
266    }
267
268    /// Revoke and remove a route. Idempotent for unknown tokens.
269    pub fn revoke(&self, token: &CapabilityToken) -> bool {
270        let mut map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
271        if let Some(binding) = map.remove(token) {
272            binding.state.store(STATE_REVOKED, Ordering::SeqCst);
273            true
274        } else {
275            false
276        }
277    }
278
279    /// Revoke every route (shutdown). Returns hex tokens that were live.
280    pub fn revoke_all(&self) -> Vec<String> {
281        let mut map = self.routes.lock().unwrap_or_else(|e| e.into_inner());
282        let mut tokens = Vec::with_capacity(map.len());
283        for (token, binding) in map.drain() {
284            binding.state.store(STATE_REVOKED, Ordering::SeqCst);
285            tokens.push(token.to_hex());
286        }
287        tokens
288    }
289
290    /// Lookup live binding by token hex (from URL).
291    pub fn get_by_hex(&self, hex: &str) -> Option<Arc<McpBinding>> {
292        let token = CapabilityToken::from_hex(hex)?;
293        self.get(&token)
294    }
295
296    /// Lookup by token.
297    pub fn get(&self, token: &CapabilityToken) -> Option<Arc<McpBinding>> {
298        self.routes
299            .lock()
300            .unwrap_or_else(|e| e.into_inner())
301            .get(token)
302            .cloned()
303    }
304
305    /// State for a token if present.
306    pub fn state_of(&self, token: &CapabilityToken) -> Option<McpBindingState> {
307        self.get(token).map(|b| b.state())
308    }
309}