agent_client_protocol_http/
server.rs1use 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}