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

agent_api/
lib.rs
