Skip to main content

mlua_swarm_server/
binding.rs

1//! Server-side implementation of the platform-neutral agent binding IF.
2//!
3//! Operator/MainAI manifests are looked up through logical role aliases.
4//! The provider returns untrusted receipts; validation and digest ownership
5//! remain in `mlua-swarm` Core.
6
7use crate::operator_ws::login::OperatorSessionEntry;
8use async_trait::async_trait;
9use mlua_swarm::{
10    AgentBindingProvider, BindOutcome, BindReceipt, BindRequest, BindingBackend,
11    BindingProviderError, ManifestBindingProvider, SessionId,
12};
13use std::collections::HashMap;
14use std::sync::Arc;
15use tokio::sync::Mutex;
16
17/// Binding provider backed by live Operator login records.
18pub struct OperatorSessionBindingProvider {
19    operator_sessions: Arc<Mutex<HashMap<SessionId, Arc<OperatorSessionEntry>>>>,
20    roles_to_sid: Arc<Mutex<HashMap<String, SessionId>>>,
21}
22
23impl OperatorSessionBindingProvider {
24    /// Bind the provider to the same session and role maps used by the
25    /// Operator REST/WebSocket login flow.
26    pub fn new(
27        operator_sessions: Arc<Mutex<HashMap<SessionId, Arc<OperatorSessionEntry>>>>,
28        roles_to_sid: Arc<Mutex<HashMap<String, SessionId>>>,
29    ) -> Self {
30        Self {
31            operator_sessions,
32            roles_to_sid,
33        }
34    }
35
36    async fn bind_operator(
37        &self,
38        request: &BindRequest,
39    ) -> Result<BindOutcome, BindingProviderError> {
40        // A WS-backed agent with no logical binding target is a Blueprint
41        // declaration error, not a transient capability gap — keep it
42        // fail-closed rather than reporting `Unbound`.
43        let target = request.binding_target.as_deref().ok_or_else(|| {
44            BindingProviderError::Provider(format!(
45                "agent '{}' uses {:?} but declares no logical binding target",
46                request.agent, request.backend
47            ))
48        })?;
49        // (a) role not joined, (b) session gone, (c) no capability_manifest:
50        // the execution environment simply has nothing to attest yet. These
51        // are `Unbound` (observed, not fatal) — the non-strict launch runs
52        // DeclarationOnly and `strict_binding` decides whether they fail.
53        // (a)/(b) would fail again at real spawn-time routing anyway, so the
54        // binding stage does not pre-gate them.
55        let Some(sid) = self.roles_to_sid.lock().await.get(target).cloned() else {
56            return Ok(BindOutcome::Unbound {
57                agent: request.agent.clone(),
58                reason: format!("no Operator session owns binding target '{target}'"),
59            });
60        };
61        let Some(entry) = self.operator_sessions.lock().await.get(&sid).cloned() else {
62            return Ok(BindOutcome::Unbound {
63                agent: request.agent.clone(),
64                reason: format!(
65                    "Operator session '{sid}' for binding target '{target}' disappeared"
66                ),
67            });
68        };
69        let Some(manifest) = entry.capability_manifest.as_ref() else {
70            return Ok(BindOutcome::Unbound {
71                agent: request.agent.clone(),
72                reason: format!(
73                    "Operator session '{sid}' for binding target '{target}' supplied no capability_manifest"
74                ),
75            });
76        };
77        // (d) manifest lacks the requested variant surfaces as `Unbound` from
78        // the delegated `ManifestBindingProvider`; a duplicate variant stays
79        // an error there. Either way the single outcome is passed straight
80        // through.
81        ManifestBindingProvider::new(manifest.clone())
82            .bind(std::slice::from_ref(request))
83            .await?
84            .pop()
85            .ok_or_else(|| {
86                BindingProviderError::Provider(format!(
87                    "Operator provider '{}' returned no outcome for agent '{}'",
88                    manifest.provider_id, request.agent
89                ))
90            })
91    }
92}
93
94#[async_trait]
95impl AgentBindingProvider for OperatorSessionBindingProvider {
96    async fn bind(
97        &self,
98        requests: &[BindRequest],
99    ) -> Result<Vec<BindOutcome>, BindingProviderError> {
100        let mut outcomes = Vec::with_capacity(requests.len());
101        for request in requests {
102            let outcome = match request.backend {
103                BindingBackend::WsOperator | BindingBackend::WsClaudeCode => {
104                    self.bind_operator(request).await?
105                }
106                // In-process AgentBlock still echoes a receipt (Core
107                // validates it); the registry-backed real attest is a future
108                // follow-up.
109                BindingBackend::AgentBlockInProcess => BindOutcome::Bound {
110                    receipt: BindReceipt {
111                        agent: request.agent.clone(),
112                        request_digest: request.request_digest.clone(),
113                        provider_id: "mse-agent-block-in-process".to_string(),
114                        provider_revision: Some(env!("CARGO_PKG_VERSION").to_string()),
115                        resolved_model: request.requested_model.clone(),
116                        effective_tools: request.requested_tools.clone(),
117                        launch_variant: None,
118                        capability_snapshot_digest: None,
119                    },
120                },
121            };
122            outcomes.push(outcome);
123        }
124        Ok(outcomes)
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use mlua_swarm::{AgentProviderCapability, AgentProviderManifest, BindingDigest};
132
133    fn request() -> BindRequest {
134        BindRequest {
135            agent: "coder".to_string(),
136            request_digest: BindingDigest::sha256("request"),
137            backend: BindingBackend::WsOperator,
138            binding_target: Some("main-ai".to_string()),
139            requested_model: Some("sonnet".to_string()),
140            requested_tools: vec!["Read".to_string()],
141            launch_variant: Some("mse-coder".to_string()),
142        }
143    }
144
145    async fn provider(manifest: Option<AgentProviderManifest>) -> OperatorSessionBindingProvider {
146        let sid = SessionId::new();
147        let entry = Arc::new(OperatorSessionEntry {
148            sid: sid.clone(),
149            token: "token".to_string(),
150            roles: vec!["main-ai".to_string()],
151            capability_manifest: manifest,
152            joined_at_secs: 0,
153            ws_session: Mutex::new(None),
154        });
155        let sessions = Arc::new(Mutex::new(HashMap::from([(sid.clone(), entry)])));
156        let roles = Arc::new(Mutex::new(HashMap::from([("main-ai".to_string(), sid)])));
157        OperatorSessionBindingProvider::new(sessions, roles)
158    }
159
160    fn expect_bound(outcome: &BindOutcome) -> &mlua_swarm::BindReceipt {
161        match outcome {
162            BindOutcome::Bound { receipt } => receipt,
163            BindOutcome::Unbound { agent, reason } => {
164                panic!("expected Bound, got Unbound({agent}): {reason}")
165            }
166        }
167    }
168
169    #[tokio::test]
170    async fn operator_manifest_resolves_to_untrusted_receipt() {
171        let manifest = AgentProviderManifest {
172            provider_id: "main-ai-self-report".to_string(),
173            provider_revision: Some("1".to_string()),
174            capabilities: vec![AgentProviderCapability {
175                launch_variant: Some("mse-coder".to_string()),
176                resolved_model: Some("claude-sonnet-4".to_string()),
177                effective_tools: vec!["Read".to_string(), "Write".to_string()],
178                capability_snapshot_digest: Some(BindingDigest::sha256("manifest")),
179            }],
180        };
181        let outcomes = provider(Some(manifest))
182            .await
183            .bind(&[request()])
184            .await
185            .unwrap();
186        assert_eq!(outcomes.len(), 1);
187        let receipt = expect_bound(&outcomes[0]);
188        assert_eq!(receipt.provider_id, "main-ai-self-report");
189        assert_eq!(receipt.request_digest, request().request_digest);
190        assert_eq!(receipt.effective_tools, ["Read", "Write"]);
191    }
192
193    #[tokio::test]
194    async fn missing_manifest_reports_unbound() {
195        let outcomes = provider(None).await.bind(&[request()]).await.unwrap();
196        assert_eq!(outcomes.len(), 1);
197        match &outcomes[0] {
198            BindOutcome::Unbound { agent, reason } => {
199                assert_eq!(agent, "coder");
200                assert!(
201                    reason.contains("supplied no capability_manifest"),
202                    "reason: {reason}"
203                );
204            }
205            BindOutcome::Bound { .. } => panic!("expected Unbound when no manifest was submitted"),
206        }
207    }
208
209    #[tokio::test]
210    async fn missing_role_reports_unbound() {
211        // A provider whose role maps are empty: the requested binding target
212        // has not joined, so the agent is Unbound (not a hard error).
213        let sessions = Arc::new(Mutex::new(HashMap::new()));
214        let roles = Arc::new(Mutex::new(HashMap::new()));
215        let provider = OperatorSessionBindingProvider::new(sessions, roles);
216        let outcomes = provider.bind(&[request()]).await.unwrap();
217        match &outcomes[0] {
218            BindOutcome::Unbound { agent, reason } => {
219                assert_eq!(agent, "coder");
220                assert!(
221                    reason.contains("no Operator session owns"),
222                    "reason: {reason}"
223                );
224            }
225            BindOutcome::Bound { .. } => panic!("expected Unbound when the role has not joined"),
226        }
227    }
228
229    #[tokio::test]
230    async fn in_process_backend_is_attested_by_server_registry() {
231        let mut request = request();
232        request.backend = BindingBackend::AgentBlockInProcess;
233        request.binding_target = None;
234        request.launch_variant = None;
235        let outcomes = provider(None).await.bind(&[request.clone()]).await.unwrap();
236        let receipt = expect_bound(&outcomes[0]);
237        assert_eq!(receipt.provider_id, "mse-agent-block-in-process");
238        assert_eq!(receipt.effective_tools, request.requested_tools);
239    }
240}