agent_client_protocol_http/
server.rs

1use std::sync::Arc;
2
3use agent_client_protocol::{Client, ConnectTo};
4use axum::{
5    Router,
6    extract::WebSocketUpgrade,
7    extract::ws::rejection::WebSocketUpgradeRejection,
8    http::{HeaderName, HeaderValue, Method, StatusCode, header, header::InvalidHeaderValue},
9    response::{IntoResponse, Response},
10    routing::{delete, get, post},
11};
12use tower_http::cors::{AllowOrigin, CorsLayer};
13
14use crate::connection::ConnectionRegistry;
15
16#[derive(Debug, Clone)]
17pub struct ServerOptions {
18    pub path: String,
19    pub cors: CorsOptions,
20    pub health_endpoint: bool,
21}
22
23impl Default for ServerOptions {
24    fn default() -> Self {
25        Self {
26            path: "/acp".to_string(),
27            cors: CorsOptions::default(),
28            health_endpoint: true,
29        }
30    }
31}
32
33#[derive(Debug, Clone, Default, PartialEq, Eq)]
34pub enum CorsOptions {
35    #[default]
36    Disabled,
37    AllowOrigins(Vec<HeaderValue>),
38    AllowAnyOrigin,
39}
40
41impl CorsOptions {
42    #[must_use]
43    pub fn disabled() -> Self {
44        Self::Disabled
45    }
46
47    #[must_use]
48    pub fn allow_any_origin() -> Self {
49        Self::AllowAnyOrigin
50    }
51
52    pub fn allow_origins<I, S>(origins: I) -> Result<Self, InvalidHeaderValue>
53    where
54        I: IntoIterator<Item = S>,
55        S: AsRef<str>,
56    {
57        origins
58            .into_iter()
59            .map(|origin| HeaderValue::from_str(origin.as_ref()))
60            .collect::<Result<Vec<_>, _>>()
61            .map(Self::AllowOrigins)
62    }
63
64    fn allow_origin_layer(&self) -> Option<AllowOrigin> {
65        match self {
66            Self::Disabled => None,
67            Self::AllowOrigins(origins) => Some(AllowOrigin::list(origins.clone())),
68            Self::AllowAnyOrigin => Some(AllowOrigin::any()),
69        }
70    }
71
72    fn allows_origin(&self, origin: Option<&HeaderValue>) -> bool {
73        let Some(origin) = origin else {
74            return true;
75        };
76        match self {
77            Self::Disabled => false,
78            Self::AllowOrigins(origins) => origins.iter().any(|allowed| allowed == origin),
79            Self::AllowAnyOrigin => true,
80        }
81    }
82}
83
84#[derive(Clone)]
85struct ServerState {
86    registry: Arc<ConnectionRegistry>,
87    cors: CorsOptions,
88}
89
90pub struct AcpHttpServer {
91    registry: Arc<ConnectionRegistry>,
92    options: ServerOptions,
93}
94
95impl std::fmt::Debug for AcpHttpServer {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        f.debug_struct("AcpHttpServer")
98            .field("options", &self.options)
99            .finish_non_exhaustive()
100    }
101}
102
103impl AcpHttpServer {
104    pub fn new<F, C>(factory: F) -> Self
105    where
106        F: Fn() -> C + Send + Sync + 'static,
107        C: ConnectTo<Client>,
108    {
109        Self {
110            registry: Arc::new(ConnectionRegistry::new(Arc::new(factory))),
111            options: ServerOptions::default(),
112        }
113    }
114
115    #[must_use]
116    pub fn with_options(mut self, options: ServerOptions) -> Self {
117        self.options = options;
118        self
119    }
120
121    pub fn into_router(self) -> Router {
122        let registry = self.registry.clone();
123        let path = self.options.path.clone();
124        let cors = self.options.cors.clone();
125        let state = ServerState {
126            registry: registry.clone(),
127            cors: cors.clone(),
128        };
129
130        let mut router = Router::new()
131            .route(
132                &path,
133                post(crate::http_server::handle_post).with_state(registry.clone()),
134            )
135            .route(&path, get(handle_get).with_state(state))
136            .route(
137                &path,
138                delete(crate::http_server::handle_delete).with_state(registry),
139            );
140
141        if self.options.health_endpoint {
142            router = router.route("/health", get(health));
143        }
144
145        if let Some(allow_origin) = cors.allow_origin_layer() {
146            router = router.layer(default_cors(allow_origin));
147        }
148
149        router
150    }
151}
152
153async fn health() -> &'static str {
154    "ok"
155}
156
157fn default_cors(allow_origin: AllowOrigin) -> CorsLayer {
158    CorsLayer::new()
159        .allow_origin(allow_origin)
160        .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::OPTIONS])
161        .allow_headers([
162            header::CONTENT_TYPE,
163            header::ACCEPT,
164            HeaderName::from_static("acp-connection-id"),
165            HeaderName::from_static("acp-session-id"),
166            header::SEC_WEBSOCKET_VERSION,
167            header::SEC_WEBSOCKET_KEY,
168            header::CONNECTION,
169            header::UPGRADE,
170        ])
171        .expose_headers([
172            HeaderName::from_static("acp-connection-id"),
173            HeaderName::from_static("acp-session-id"),
174        ])
175}
176
177async fn handle_get(
178    ws_upgrade: Result<WebSocketUpgrade, WebSocketUpgradeRejection>,
179    axum::extract::State(state): axum::extract::State<ServerState>,
180    request: axum::http::Request<axum::body::Body>,
181) -> Response {
182    match ws_upgrade {
183        Ok(ws) => {
184            if !state
185                .cors
186                .allows_origin(request.headers().get(header::ORIGIN))
187            {
188                return (StatusCode::FORBIDDEN, "WebSocket origin not allowed").into_response();
189            }
190            crate::websocket_server::handle_ws_upgrade(state.registry, ws)
191        }
192        Err(_) => crate::http_server::handle_get(state.registry, request).await,
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use axum::body::Body;
200    use tower::{Layer as _, ServiceExt as _, service_fn};
201
202    #[test]
203    fn cors_is_disabled_by_default() {
204        assert_eq!(ServerOptions::default().cors, CorsOptions::Disabled);
205    }
206
207    #[test]
208    fn disabled_cors_rejects_browser_origin_for_websockets() {
209        let origin = HeaderValue::from_static("http://localhost:5173");
210
211        assert!(CorsOptions::disabled().allows_origin(None));
212        assert!(!CorsOptions::disabled().allows_origin(Some(&origin)));
213    }
214
215    #[test]
216    fn cors_allowlist_matches_configured_origins() {
217        let allowed = HeaderValue::from_static("http://localhost:5173");
218        let denied = HeaderValue::from_static("http://localhost:3000");
219        let cors = CorsOptions::allow_origins(["http://localhost:5173"]).unwrap();
220
221        assert!(cors.allows_origin(None));
222        assert!(cors.allows_origin(Some(&allowed)));
223        assert!(!cors.allows_origin(Some(&denied)));
224    }
225
226    #[test]
227    fn explicit_allow_any_origin_accepts_browser_origins() {
228        let origin = HeaderValue::from_static("https://example.com");
229
230        assert!(CorsOptions::allow_any_origin().allows_origin(Some(&origin)));
231    }
232
233    #[tokio::test]
234    async fn allow_any_origin_uses_wildcard_cors_header() {
235        let response = default_cors(
236            CorsOptions::allow_any_origin()
237                .allow_origin_layer()
238                .expect("CORS layer"),
239        )
240        .layer(service_fn(|_: axum::http::Request<Body>| async {
241            Ok::<_, std::convert::Infallible>(Response::new(Body::empty()))
242        }))
243        .oneshot(
244            axum::http::Request::builder()
245                .header(header::ORIGIN, "https://example.com")
246                .body(Body::empty())
247                .unwrap(),
248        )
249        .await
250        .unwrap();
251
252        assert_eq!(
253            response.headers().get(header::ACCESS_CONTROL_ALLOW_ORIGIN),
254            Some(&HeaderValue::from_static("*"))
255        );
256        assert!(response.headers().get(header::VARY).is_none());
257    }
258
259    #[tokio::test]
260    async fn allowlisted_origins_vary_by_origin() {
261        let response = default_cors(
262            CorsOptions::allow_origins(["https://example.com"])
263                .unwrap()
264                .allow_origin_layer()
265                .expect("CORS layer"),
266        )
267        .layer(service_fn(|_: axum::http::Request<Body>| async {
268            Ok::<_, std::convert::Infallible>(Response::new(Body::empty()))
269        }))
270        .oneshot(
271            axum::http::Request::builder()
272                .header(header::ORIGIN, "https://example.com")
273                .body(Body::empty())
274                .unwrap(),
275        )
276        .await
277        .unwrap();
278
279        assert_eq!(
280            response.headers().get(header::ACCESS_CONTROL_ALLOW_ORIGIN),
281            Some(&HeaderValue::from_static("https://example.com"))
282        );
283        assert_eq!(
284            response.headers().get(header::VARY),
285            Some(&HeaderValue::from_static("origin"))
286        );
287    }
288}
agent_client_protocol_http/server.rs

agent_client_protocol_http/
server.rs
