agent_api/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::future::Future;
5use std::path::PathBuf;
6use std::pin::Pin;
7use std::process::ExitStatus;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::sync::Arc;
10use std::time::Duration;
11
12use futures_core::Stream;
13
14use crate::mcp::{
15    normalize_mcp_add_request, normalize_mcp_get_request, normalize_mcp_remove_request,
16    AgentWrapperMcpAddRequest, AgentWrapperMcpCommandOutput, AgentWrapperMcpGetRequest,
17    AgentWrapperMcpListRequest, AgentWrapperMcpRemoveRequest, CAPABILITY_MCP_ADD_V1,
18    CAPABILITY_MCP_GET_V1, CAPABILITY_MCP_LIST_V1, CAPABILITY_MCP_REMOVE_V1,
19};
20
21#[cfg(any(feature = "codex", feature = "claude_code", feature = "opencode"))]
22mod bounds;
23
24#[cfg(any(feature = "codex", feature = "claude_code", feature = "opencode"))]
25mod run_handle_gate;
26
27#[cfg(any(feature = "codex", feature = "claude_code", feature = "opencode"))]
28mod backend_harness;
29
30pub mod backends;
31pub mod mcp;
32
33const CAPABILITY_CONTROL_CANCEL_V1: &str = "agent_api.control.cancel.v1";
34#[allow(dead_code)]
35#[cfg(any(feature = "codex", feature = "claude_code", feature = "opencode", test))]
36pub(crate) const EXT_AGENT_API_CONFIG_MODEL_V1: &str = "agent_api.config.model.v1";
37
38#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
39pub struct AgentWrapperKind(String);
40
41impl AgentWrapperKind {
42    /// Creates an agent kind from a string.
43    ///
44    /// The value MUST follow `capabilities-schema-spec.md` naming rules.
45    pub fn new(value: impl Into<String>) -> Result<Self, AgentWrapperError> {
46        let value = value.into();
47        validate_agent_kind(&value)?;
48        Ok(Self(value))
49    }
50
51    /// Returns the canonical string id.
52    pub fn as_str(&self) -> &str {
53        &self.0
54    }
55}
56
57#[derive(Clone, Debug, Default, Eq, PartialEq)]
58pub struct AgentWrapperCapabilities {
59    /// Set of namespaced capability ids (see `capabilities-schema-spec.md`).
60    pub ids: BTreeSet<String>,
61}
62
63impl AgentWrapperCapabilities {
64    pub fn contains(&self, capability_id: &str) -> bool {
65        self.ids.contains(capability_id)
66    }
67}
68
69#[derive(Clone, Debug, Eq, PartialEq)]
70pub enum AgentWrapperEventKind {
71    TextOutput,
72    ToolCall,
73    ToolResult,
74    Status,
75    Error,
76    Unknown,
77}
78
79#[derive(Clone, Debug, Eq, PartialEq)]
80pub struct AgentWrapperEvent {
81    pub agent_kind: AgentWrapperKind,
82    pub kind: AgentWrapperEventKind,
83    pub channel: Option<String>,
84    /// Stable payload for `TextOutput` events.
85    pub text: Option<String>,
86    /// Stable payload for `Status` and `Error` events.
87    pub message: Option<String>,
88    pub data: Option<serde_json::Value>,
89}
90
91#[derive(Clone, Debug, Default)]
92pub struct AgentWrapperRunRequest {
93    pub prompt: String,
94    pub working_dir: Option<PathBuf>,
95    pub timeout: Option<Duration>,
96    pub env: BTreeMap<String, String>,
97    /// Extension options are namespaced keys with JSON values.
98    pub extensions: BTreeMap<String, serde_json::Value>,
99}
100
101pub type DynAgentWrapperEventStream = Pin<Box<dyn Stream<Item = AgentWrapperEvent> + Send>>;
102pub type DynAgentWrapperCompletion =
103    Pin<Box<dyn Future<Output = Result<AgentWrapperCompletion, AgentWrapperError>> + Send>>;
104
105pub struct AgentWrapperRunHandle {
106    pub events: DynAgentWrapperEventStream,
107    pub completion: DynAgentWrapperCompletion,
108}
109
110impl std::fmt::Debug for AgentWrapperRunHandle {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        f.debug_struct("AgentWrapperRunHandle")
113            .field("events", &"<stream>")
114            .field("completion", &"<future>")
115            .finish()
116    }
117}
118
119struct AgentWrapperCancelInner {
120    called: AtomicBool,
121    cancel: Box<dyn Fn() + Send + Sync + 'static>,
122}
123
124#[derive(Clone)]
125pub struct AgentWrapperCancelHandle {
126    inner: Arc<AgentWrapperCancelInner>,
127}
128
129impl std::fmt::Debug for AgentWrapperCancelHandle {
130    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131        f.debug_struct("AgentWrapperCancelHandle")
132            .finish_non_exhaustive()
133    }
134}
135
136impl AgentWrapperCancelHandle {
137    #[allow(dead_code)]
138    pub(crate) fn new(cancel: impl Fn() + Send + Sync + 'static) -> Self {
139        Self {
140            inner: Arc::new(AgentWrapperCancelInner {
141                called: AtomicBool::new(false),
142                cancel: Box::new(cancel),
143            }),
144        }
145    }
146
147    /// Requests best-effort cancellation of the underlying backend process.
148    ///
149    /// This method MUST be idempotent.
150    ///
151    /// If cancellation is requested before `AgentWrapperRunHandle.completion` resolves, the completion
152    /// MUST resolve to `Err(AgentWrapperError::Backend { message: "cancelled" })`.
153    ///
154    /// Canonical semantics: `docs/specs/unified-agent-api/run-protocol-spec.md` ("Explicit
155    /// cancellation semantics").
156    pub fn cancel(&self) {
157        if self.inner.called.swap(true, Ordering::SeqCst) {
158            return;
159        }
160
161        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
162            (self.inner.cancel)();
163        }));
164    }
165}
166
167#[derive(Debug)]
168pub struct AgentWrapperRunControl {
169    pub handle: AgentWrapperRunHandle,
170    pub cancel: AgentWrapperCancelHandle,
171}
172
173#[derive(Clone, Debug)]
174pub struct AgentWrapperCompletion {
175    pub status: ExitStatus,
176    /// A backend may populate `final_text` when it can deterministically extract it.
177    pub final_text: Option<String>,
178    /// Optional backend-specific completion payload.
179    ///
180    /// This payload MUST obey the bounds and enforcement behavior defined in
181    /// `event-envelope-schema-spec.md` (see "Completion payload bounds").
182    pub data: Option<serde_json::Value>,
183}
184
185#[derive(Clone, Debug)]
186pub struct AgentWrapperRunResult {
187    pub completion: AgentWrapperCompletion,
188}
189
190#[derive(Debug, thiserror::Error)]
191pub enum AgentWrapperError {
192    #[error("unknown backend: {agent_kind}")]
193    UnknownBackend { agent_kind: String },
194    #[error("unsupported capability for {agent_kind}: {capability}")]
195    UnsupportedCapability {
196        agent_kind: String,
197        capability: String,
198    },
199    #[error("invalid agent kind: {message}")]
200    InvalidAgentKind { message: String },
201    #[error("invalid request: {message}")]
202    InvalidRequest { message: String },
203    #[error("backend error: {message}")]
204    Backend { message: String },
205}
206
207pub trait AgentWrapperBackend: Send + Sync {
208    fn kind(&self) -> AgentWrapperKind;
209    fn capabilities(&self) -> AgentWrapperCapabilities;
210
211    /// Starts a run and returns a handle producing events and a completion result.
212    ///
213    /// Backends MUST enforce capability gating per `run-protocol-spec.md`.
214    fn run(
215        &self,
216        request: AgentWrapperRunRequest,
217    ) -> Pin<Box<dyn Future<Output = Result<AgentWrapperRunHandle, AgentWrapperError>> + Send + '_>>;
218
219    /// Starts a run and returns a handle plus an explicit cancellation handle.
220    ///
221    /// Backends that do not advertise `agent_api.control.cancel.v1` MUST return:
222    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability: "agent_api.control.cancel.v1" }`,
223    /// where `agent_kind == self.kind().as_str().to_string()`.
224    fn run_control(
225        &self,
226        _request: AgentWrapperRunRequest,
227    ) -> Pin<Box<dyn Future<Output = Result<AgentWrapperRunControl, AgentWrapperError>> + Send + '_>>
228    {
229        let agent_kind = self.kind().as_str().to_string();
230        Box::pin(async move {
231            Err(AgentWrapperError::UnsupportedCapability {
232                agent_kind,
233                capability: CAPABILITY_CONTROL_CANCEL_V1.to_string(),
234            })
235        })
236    }
237
238    /// Runs `mcp list` and returns bounded command output.
239    ///
240    /// Backends that do not advertise `agent_api.tools.mcp.list.v1` MUST return:
241    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability: "agent_api.tools.mcp.list.v1" }`,
242    /// where `agent_kind == self.kind().as_str().to_string()`.
243    fn mcp_list(
244        &self,
245        _request: AgentWrapperMcpListRequest,
246    ) -> Pin<
247        Box<
248            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
249                + Send
250                + '_,
251        >,
252    > {
253        let agent_kind = self.kind().as_str().to_string();
254        Box::pin(async move {
255            Err(AgentWrapperError::UnsupportedCapability {
256                agent_kind,
257                capability: CAPABILITY_MCP_LIST_V1.to_string(),
258            })
259        })
260    }
261
262    /// Runs `mcp get` and returns bounded command output.
263    ///
264    /// Backends that do not advertise `agent_api.tools.mcp.get.v1` MUST return:
265    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability: "agent_api.tools.mcp.get.v1" }`,
266    /// where `agent_kind == self.kind().as_str().to_string()`.
267    fn mcp_get(
268        &self,
269        _request: AgentWrapperMcpGetRequest,
270    ) -> Pin<
271        Box<
272            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
273                + Send
274                + '_,
275        >,
276    > {
277        let agent_kind = self.kind().as_str().to_string();
278        Box::pin(async move {
279            Err(AgentWrapperError::UnsupportedCapability {
280                agent_kind,
281                capability: CAPABILITY_MCP_GET_V1.to_string(),
282            })
283        })
284    }
285
286    /// Runs `mcp add` and returns bounded command output.
287    ///
288    /// Backends that do not advertise `agent_api.tools.mcp.add.v1` MUST return:
289    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability: "agent_api.tools.mcp.add.v1" }`,
290    /// where `agent_kind == self.kind().as_str().to_string()`.
291    fn mcp_add(
292        &self,
293        _request: AgentWrapperMcpAddRequest,
294    ) -> Pin<
295        Box<
296            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
297                + Send
298                + '_,
299        >,
300    > {
301        let agent_kind = self.kind().as_str().to_string();
302        Box::pin(async move {
303            Err(AgentWrapperError::UnsupportedCapability {
304                agent_kind,
305                capability: CAPABILITY_MCP_ADD_V1.to_string(),
306            })
307        })
308    }
309
310    /// Runs `mcp remove` and returns bounded command output.
311    ///
312    /// Backends that do not advertise `agent_api.tools.mcp.remove.v1` MUST return:
313    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability: "agent_api.tools.mcp.remove.v1" }`,
314    /// where `agent_kind == self.kind().as_str().to_string()`.
315    fn mcp_remove(
316        &self,
317        _request: AgentWrapperMcpRemoveRequest,
318    ) -> Pin<
319        Box<
320            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
321                + Send
322                + '_,
323        >,
324    > {
325        let agent_kind = self.kind().as_str().to_string();
326        Box::pin(async move {
327            Err(AgentWrapperError::UnsupportedCapability {
328                agent_kind,
329                capability: CAPABILITY_MCP_REMOVE_V1.to_string(),
330            })
331        })
332    }
333}
334
335#[derive(Clone, Default)]
336pub struct AgentWrapperGateway {
337    backends: BTreeMap<AgentWrapperKind, Arc<dyn AgentWrapperBackend>>,
338}
339
340impl AgentWrapperGateway {
341    pub fn new() -> Self {
342        Self::default()
343    }
344
345    /// Registers a backend.
346    ///
347    /// If a backend with the same `AgentWrapperKind` is already registered, this MUST return an error.
348    pub fn register(
349        &mut self,
350        backend: Arc<dyn AgentWrapperBackend>,
351    ) -> Result<(), AgentWrapperError> {
352        let kind = backend.kind();
353        if self.backends.contains_key(&kind) {
354            return Err(AgentWrapperError::InvalidRequest {
355                message: format!("backend already registered: {}", kind.as_str()),
356            });
357        }
358        self.backends.insert(kind, backend);
359        Ok(())
360    }
361
362    /// Resolves a backend by `AgentWrapperKind`.
363    pub fn backend(&self, agent_kind: &AgentWrapperKind) -> Option<Arc<dyn AgentWrapperBackend>> {
364        self.backends.get(agent_kind).cloned()
365    }
366
367    /// Convenience entrypoint: resolves a backend and starts a run.
368    ///
369    /// This MUST return `AgentWrapperError::UnknownBackend` when no backend is registered for `agent_kind`.
370    pub fn run(
371        &self,
372        agent_kind: &AgentWrapperKind,
373        request: AgentWrapperRunRequest,
374    ) -> Pin<Box<dyn Future<Output = Result<AgentWrapperRunHandle, AgentWrapperError>> + Send + '_>>
375    {
376        let backend = self.backends.get(agent_kind).cloned();
377        let agent_kind = agent_kind.as_str().to_string();
378        Box::pin(async move {
379            let backend = backend.ok_or(AgentWrapperError::UnknownBackend { agent_kind })?;
380            backend.run(request).await
381        })
382    }
383
384    /// Starts a run and returns a control object including an explicit cancellation handle.
385    ///
386    /// This MUST return `AgentWrapperError::UnknownBackend { agent_kind }` when no backend is
387    /// registered for the requested `agent_kind`, where
388    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
389    ///
390    /// If the resolved backend does not advertise `agent_api.control.cancel.v1`, this MUST return:
391    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability: "agent_api.control.cancel.v1" }`,
392    /// where `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
393    ///
394    /// Cancellation is best-effort and is defined by
395    /// `docs/specs/unified-agent-api/run-protocol-spec.md`, including the pinned `"cancelled"`
396    /// completion outcome.
397    pub fn run_control(
398        &self,
399        agent_kind: &AgentWrapperKind,
400        request: AgentWrapperRunRequest,
401    ) -> Pin<Box<dyn Future<Output = Result<AgentWrapperRunControl, AgentWrapperError>> + Send + '_>>
402    {
403        let backend = self.backends.get(agent_kind).cloned();
404        let agent_kind = agent_kind.as_str().to_string();
405        Box::pin(async move {
406            let backend = backend.ok_or(AgentWrapperError::UnknownBackend {
407                agent_kind: agent_kind.clone(),
408            })?;
409
410            if !backend
411                .capabilities()
412                .contains(CAPABILITY_CONTROL_CANCEL_V1)
413            {
414                return Err(AgentWrapperError::UnsupportedCapability {
415                    agent_kind,
416                    capability: CAPABILITY_CONTROL_CANCEL_V1.to_string(),
417                });
418            }
419
420            backend.run_control(request).await
421        })
422    }
423
424    /// Resolves a backend and runs `mcp list`.
425    ///
426    /// This MUST return `AgentWrapperError::UnknownBackend { agent_kind }` when no backend is
427    /// registered for the requested `agent_kind`, where
428    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
429    ///
430    /// If the resolved backend does not advertise `agent_api.tools.mcp.list.v1`, this MUST
431    /// return `AgentWrapperError::UnsupportedCapability { agent_kind, capability:
432    /// "agent_api.tools.mcp.list.v1" }`, where
433    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
434    pub fn mcp_list(
435        &self,
436        agent_kind: &AgentWrapperKind,
437        request: AgentWrapperMcpListRequest,
438    ) -> Pin<
439        Box<
440            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
441                + Send
442                + '_,
443        >,
444    > {
445        let backend = self.backends.get(agent_kind).cloned();
446        let agent_kind = agent_kind.as_str().to_string();
447        Box::pin(async move {
448            let backend = backend.ok_or(AgentWrapperError::UnknownBackend {
449                agent_kind: agent_kind.clone(),
450            })?;
451
452            if !backend.capabilities().contains(CAPABILITY_MCP_LIST_V1) {
453                return Err(AgentWrapperError::UnsupportedCapability {
454                    agent_kind,
455                    capability: CAPABILITY_MCP_LIST_V1.to_string(),
456                });
457            }
458
459            backend.mcp_list(request).await
460        })
461    }
462
463    /// Resolves a backend and runs `mcp get`.
464    ///
465    /// This MUST return `AgentWrapperError::UnknownBackend { agent_kind }` when no backend is
466    /// registered for the requested `agent_kind`, where
467    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
468    ///
469    /// If the resolved backend does not advertise `agent_api.tools.mcp.get.v1`, this MUST return
470    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability:
471    /// "agent_api.tools.mcp.get.v1" }`, where
472    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
473    pub fn mcp_get(
474        &self,
475        agent_kind: &AgentWrapperKind,
476        request: AgentWrapperMcpGetRequest,
477    ) -> Pin<
478        Box<
479            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
480                + Send
481                + '_,
482        >,
483    > {
484        let backend = self.backends.get(agent_kind).cloned();
485        let agent_kind = agent_kind.as_str().to_string();
486        Box::pin(async move {
487            let backend = backend.ok_or(AgentWrapperError::UnknownBackend {
488                agent_kind: agent_kind.clone(),
489            })?;
490
491            if !backend.capabilities().contains(CAPABILITY_MCP_GET_V1) {
492                return Err(AgentWrapperError::UnsupportedCapability {
493                    agent_kind,
494                    capability: CAPABILITY_MCP_GET_V1.to_string(),
495                });
496            }
497
498            let request = normalize_mcp_get_request(request)?;
499            backend.mcp_get(request).await
500        })
501    }
502
503    /// Resolves a backend and runs `mcp add`.
504    ///
505    /// This MUST return `AgentWrapperError::UnknownBackend { agent_kind }` when no backend is
506    /// registered for the requested `agent_kind`, where
507    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
508    ///
509    /// If the resolved backend does not advertise `agent_api.tools.mcp.add.v1`, this MUST return
510    /// `AgentWrapperError::UnsupportedCapability { agent_kind, capability:
511    /// "agent_api.tools.mcp.add.v1" }`, where
512    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
513    pub fn mcp_add(
514        &self,
515        agent_kind: &AgentWrapperKind,
516        request: AgentWrapperMcpAddRequest,
517    ) -> Pin<
518        Box<
519            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
520                + Send
521                + '_,
522        >,
523    > {
524        let backend = self.backends.get(agent_kind).cloned();
525        let agent_kind = agent_kind.as_str().to_string();
526        Box::pin(async move {
527            let backend = backend.ok_or(AgentWrapperError::UnknownBackend {
528                agent_kind: agent_kind.clone(),
529            })?;
530
531            if !backend.capabilities().contains(CAPABILITY_MCP_ADD_V1) {
532                return Err(AgentWrapperError::UnsupportedCapability {
533                    agent_kind,
534                    capability: CAPABILITY_MCP_ADD_V1.to_string(),
535                });
536            }
537
538            let request = normalize_mcp_add_request(request)?;
539            backend.mcp_add(request).await
540        })
541    }
542
543    /// Resolves a backend and runs `mcp remove`.
544    ///
545    /// This MUST return `AgentWrapperError::UnknownBackend { agent_kind }` when no backend is
546    /// registered for the requested `agent_kind`, where
547    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
548    ///
549    /// If the resolved backend does not advertise `agent_api.tools.mcp.remove.v1`, this MUST
550    /// return `AgentWrapperError::UnsupportedCapability { agent_kind, capability:
551    /// "agent_api.tools.mcp.remove.v1" }`, where
552    /// `agent_kind == <requested AgentWrapperKind>.as_str().to_string()`.
553    pub fn mcp_remove(
554        &self,
555        agent_kind: &AgentWrapperKind,
556        request: AgentWrapperMcpRemoveRequest,
557    ) -> Pin<
558        Box<
559            dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
560                + Send
561                + '_,
562        >,
563    > {
564        let backend = self.backends.get(agent_kind).cloned();
565        let agent_kind = agent_kind.as_str().to_string();
566        Box::pin(async move {
567            let backend = backend.ok_or(AgentWrapperError::UnknownBackend {
568                agent_kind: agent_kind.clone(),
569            })?;
570
571            if !backend.capabilities().contains(CAPABILITY_MCP_REMOVE_V1) {
572                return Err(AgentWrapperError::UnsupportedCapability {
573                    agent_kind,
574                    capability: CAPABILITY_MCP_REMOVE_V1.to_string(),
575                });
576            }
577
578            let request = normalize_mcp_remove_request(request)?;
579            backend.mcp_remove(request).await
580        })
581    }
582}
583
584fn validate_agent_kind(value: &str) -> Result<(), AgentWrapperError> {
585    if value.is_empty() {
586        return Err(AgentWrapperError::InvalidAgentKind {
587            message: "agent kind must not be empty".to_string(),
588        });
589    }
590    let mut chars = value.chars();
591    let first = chars.next().unwrap_or_default();
592    if !first.is_ascii_lowercase() {
593        return Err(AgentWrapperError::InvalidAgentKind {
594            message: "agent kind must start with a lowercase ASCII letter".to_string(),
595        });
596    }
597    for ch in chars {
598        if !(ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '_') {
599            return Err(AgentWrapperError::InvalidAgentKind {
600                message: "agent kind must match ^[a-z][a-z0-9_]*$".to_string(),
601            });
602        }
603    }
604    Ok(())
605}
606
607#[cfg(test)]
608mod tests {
609    use std::collections::BTreeSet;
610    use std::future::Future;
611    use std::pin::Pin;
612    use std::sync::atomic::{AtomicUsize, Ordering};
613    use std::sync::{Arc, Mutex};
614    use std::task::{Context, Poll, Wake, Waker};
615
616    use super::*;
617    use crate::mcp::{
618        AgentWrapperMcpAddRequest, AgentWrapperMcpAddTransport, AgentWrapperMcpCommandContext,
619        AgentWrapperMcpGetRequest, CAPABILITY_MCP_ADD_V1, CAPABILITY_MCP_GET_V1,
620    };
621
622    #[test]
623    fn cancel_handle_is_idempotent() {
624        let calls = Arc::new(AtomicUsize::new(0));
625        let calls_for_cancel = Arc::clone(&calls);
626        let cancel = AgentWrapperCancelHandle::new(move || {
627            calls_for_cancel.fetch_add(1, Ordering::SeqCst);
628        });
629
630        cancel.cancel();
631        cancel.cancel();
632        cancel.cancel();
633
634        assert_eq!(calls.load(Ordering::SeqCst), 1);
635    }
636
637    #[test]
638    fn mcp_add_returns_unknown_backend_before_validation() {
639        let gateway = AgentWrapperGateway::new();
640        let agent_kind = AgentWrapperKind::new("codex").expect("valid kind");
641        let secret = "SECRET_BACKEND_VALUE";
642
643        let err = block_on_immediate(gateway.mcp_add(
644            &agent_kind,
645            AgentWrapperMcpAddRequest {
646                name: format!("  {secret}  "),
647                transport: AgentWrapperMcpAddTransport::Url {
648                    url: format!("relative/{secret}"),
649                    bearer_token_env_var: None,
650                },
651                context: AgentWrapperMcpCommandContext::default(),
652            },
653        ))
654        .expect_err("unknown backend should fail first");
655
656        match err {
657            AgentWrapperError::UnknownBackend { agent_kind } => {
658                assert_eq!(agent_kind, "codex");
659            }
660            other => panic!("expected UnknownBackend, got {other:?}"),
661        }
662    }
663
664    #[test]
665    fn mcp_add_returns_unsupported_capability_before_validation_and_without_hook_call() {
666        let backend = Arc::new(TestBackend::new(BTreeSet::new()));
667        let mut gateway = AgentWrapperGateway::new();
668        gateway.register(backend.clone()).expect("register backend");
669        let agent_kind = backend.kind();
670        let secret = "SECRET_UNSUPPORTED_VALUE";
671
672        let err = block_on_immediate(gateway.mcp_add(
673            &agent_kind,
674            AgentWrapperMcpAddRequest {
675                name: format!("  {secret}  "),
676                transport: AgentWrapperMcpAddTransport::Url {
677                    url: format!("relative/{secret}"),
678                    bearer_token_env_var: None,
679                },
680                context: AgentWrapperMcpCommandContext::default(),
681            },
682        ))
683        .expect_err("unsupported capability should fail before validation");
684
685        match err {
686            AgentWrapperError::UnsupportedCapability {
687                agent_kind,
688                capability,
689            } => {
690                assert_eq!(agent_kind, "test_backend");
691                assert_eq!(capability, CAPABILITY_MCP_ADD_V1);
692            }
693            other => panic!("expected UnsupportedCapability, got {other:?}"),
694        }
695
696        assert_eq!(backend.mcp_add_calls.load(Ordering::SeqCst), 0);
697    }
698
699    #[test]
700    fn mcp_add_returns_invalid_request_before_hook_when_capability_is_advertised() {
701        let backend = Arc::new(TestBackend::new(BTreeSet::from([
702            CAPABILITY_MCP_ADD_V1.to_string()
703        ])));
704        let mut gateway = AgentWrapperGateway::new();
705        gateway.register(backend.clone()).expect("register backend");
706        let agent_kind = backend.kind();
707        let secret = "SECRET_INVALID_VALUE";
708
709        let err = block_on_immediate(gateway.mcp_add(
710            &agent_kind,
711            AgentWrapperMcpAddRequest {
712                name: "  example  ".to_string(),
713                transport: AgentWrapperMcpAddTransport::Url {
714                    url: format!("https:{secret}.example.com"),
715                    bearer_token_env_var: None,
716                },
717                context: AgentWrapperMcpCommandContext::default(),
718            },
719        ))
720        .expect_err("invalid request should fail before hook");
721
722        match err {
723            AgentWrapperError::InvalidRequest { message } => {
724                assert_eq!(message, "mcp add url must be an absolute http or https URL");
725                assert!(!message.contains(secret));
726            }
727            other => panic!("expected InvalidRequest, got {other:?}"),
728        }
729
730        assert_eq!(backend.mcp_add_calls.load(Ordering::SeqCst), 0);
731    }
732
733    #[test]
734    fn mcp_get_passes_normalized_request_to_hook() {
735        let backend = Arc::new(TestBackend::new(BTreeSet::from([
736            CAPABILITY_MCP_GET_V1.to_string()
737        ])));
738        let mut gateway = AgentWrapperGateway::new();
739        gateway.register(backend.clone()).expect("register backend");
740        let agent_kind = backend.kind();
741
742        let output = block_on_immediate(gateway.mcp_get(
743            &agent_kind,
744            AgentWrapperMcpGetRequest {
745                name: "  normalized-name  ".to_string(),
746                context: AgentWrapperMcpCommandContext::default(),
747            },
748        ))
749        .expect("normalized request should reach hook");
750
751        assert_eq!(output.stdout, "ok");
752        assert_eq!(backend.mcp_get_calls.load(Ordering::SeqCst), 1);
753        assert_eq!(
754            backend
755                .last_get_name
756                .lock()
757                .expect("get name mutex")
758                .as_deref(),
759            Some("normalized-name")
760        );
761    }
762
763    struct NoopWake;
764
765    impl Wake for NoopWake {
766        fn wake(self: Arc<Self>) {}
767    }
768
769    fn block_on_immediate<F>(future: F) -> F::Output
770    where
771        F: Future,
772    {
773        let waker = Waker::from(Arc::new(NoopWake));
774        let mut future = Box::pin(future);
775        let mut context = Context::from_waker(&waker);
776
777        loop {
778            match future.as_mut().poll(&mut context) {
779                Poll::Ready(output) => return output,
780                Poll::Pending => std::thread::yield_now(),
781            }
782        }
783    }
784
785    fn success_exit_status() -> std::process::ExitStatus {
786        #[cfg(unix)]
787        {
788            use std::os::unix::process::ExitStatusExt;
789            std::process::ExitStatus::from_raw(0)
790        }
791        #[cfg(windows)]
792        {
793            use std::os::windows::process::ExitStatusExt;
794            std::process::ExitStatus::from_raw(0)
795        }
796    }
797
798    fn success_output() -> AgentWrapperMcpCommandOutput {
799        AgentWrapperMcpCommandOutput {
800            status: success_exit_status(),
801            stdout: "ok".to_string(),
802            stderr: String::new(),
803            stdout_truncated: false,
804            stderr_truncated: false,
805        }
806    }
807
808    struct TestBackend {
809        capabilities: AgentWrapperCapabilities,
810        mcp_add_calls: AtomicUsize,
811        mcp_get_calls: AtomicUsize,
812        last_get_name: Mutex<Option<String>>,
813    }
814
815    impl TestBackend {
816        fn new(capabilities: BTreeSet<String>) -> Self {
817            Self {
818                capabilities: AgentWrapperCapabilities { ids: capabilities },
819                mcp_add_calls: AtomicUsize::new(0),
820                mcp_get_calls: AtomicUsize::new(0),
821                last_get_name: Mutex::new(None),
822            }
823        }
824    }
825
826    impl AgentWrapperBackend for TestBackend {
827        fn kind(&self) -> AgentWrapperKind {
828            AgentWrapperKind("test_backend".to_string())
829        }
830
831        fn capabilities(&self) -> AgentWrapperCapabilities {
832            self.capabilities.clone()
833        }
834
835        fn run(
836            &self,
837            _request: AgentWrapperRunRequest,
838        ) -> Pin<
839            Box<dyn Future<Output = Result<AgentWrapperRunHandle, AgentWrapperError>> + Send + '_>,
840        > {
841            Box::pin(async {
842                Err(AgentWrapperError::Backend {
843                    message: "run not used in tests".to_string(),
844                })
845            })
846        }
847
848        fn mcp_get(
849            &self,
850            request: AgentWrapperMcpGetRequest,
851        ) -> Pin<
852            Box<
853                dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
854                    + Send
855                    + '_,
856            >,
857        > {
858            self.mcp_get_calls.fetch_add(1, Ordering::SeqCst);
859            *self.last_get_name.lock().expect("last get name mutex") = Some(request.name);
860            Box::pin(async move { Ok(success_output()) })
861        }
862
863        fn mcp_add(
864            &self,
865            _request: AgentWrapperMcpAddRequest,
866        ) -> Pin<
867            Box<
868                dyn Future<Output = Result<AgentWrapperMcpCommandOutput, AgentWrapperError>>
869                    + Send
870                    + '_,
871            >,
872        > {
873            self.mcp_add_calls.fetch_add(1, Ordering::SeqCst);
874            Box::pin(async move { Ok(success_output()) })
875        }
876    }
877}
agent_api/lib.rs

agent_api/
lib.rs
