Skip to main content

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::mirror_request()),
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
200    #[test]
201    fn cors_is_disabled_by_default() {
202        assert_eq!(ServerOptions::default().cors, CorsOptions::Disabled);
203    }
204
205    #[test]
206    fn disabled_cors_rejects_browser_origin_for_websockets() {
207        let origin = HeaderValue::from_static("http://localhost:5173");
208
209        assert!(CorsOptions::disabled().allows_origin(None));
210        assert!(!CorsOptions::disabled().allows_origin(Some(&origin)));
211    }
212
213    #[test]
214    fn cors_allowlist_matches_configured_origins() {
215        let allowed = HeaderValue::from_static("http://localhost:5173");
216        let denied = HeaderValue::from_static("http://localhost:3000");
217        let cors = CorsOptions::allow_origins(["http://localhost:5173"]).unwrap();
218
219        assert!(cors.allows_origin(None));
220        assert!(cors.allows_origin(Some(&allowed)));
221        assert!(!cors.allows_origin(Some(&denied)));
222    }
223
224    #[test]
225    fn explicit_allow_any_origin_accepts_browser_origins() {
226        let origin = HeaderValue::from_static("https://example.com");
227
228        assert!(CorsOptions::allow_any_origin().allows_origin(Some(&origin)));
229    }
230}