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

agent_api/
lib.rs
