codex_helper_core/proxy/
router_setup.rs

1use axum::Json;
2use axum::Router;
3use axum::extract::ws::WebSocketUpgrade;
4use axum::extract::ws::rejection::WebSocketUpgradeRejection;
5use axum::http::{HeaderMap, Uri};
6use axum::middleware;
7use axum::routing::MethodFilter;
8use axum::routing::{any, get, on};
9
10use super::ProxyService;
11use super::admin::{ProxyAdminDiscovery, reject_admin_paths_from_proxy, require_admin_path_only};
12use super::control_plane_routes::control_plane_routes;
13use super::handle_proxy;
14use super::responses_websocket::handle_responses_websocket;
15
16pub fn router(proxy: ProxyService) -> Router {
17    // In axum 0.8, wildcard segments use `/{*path}` (equivalent to `/*path` from axum 0.7).
18    let proxy_routes = proxy_only_router(proxy.clone());
19
20    Router::new()
21        .merge(control_plane_routes(proxy))
22        .merge(proxy_routes)
23}
24
25pub fn proxy_only_router(proxy: ProxyService) -> Router {
26    proxy_only_router_with_admin_base_url(proxy, None)
27}
28
29pub fn proxy_only_router_with_admin_base_url(
30    proxy: ProxyService,
31    admin_base_url: Option<String>,
32) -> Router {
33    let service_name = proxy.service_name;
34    let discovery = admin_base_url.map(|admin_base_url| {
35        Json(ProxyAdminDiscovery {
36            api_version: 1,
37            service_name,
38            admin_base_url,
39        })
40    });
41
42    let mut router = Router::new();
43    if let Some(discovery) = discovery {
44        router = router.route(
45            "/.well-known/codex-helper-admin",
46            get(move || {
47                let discovery = discovery.clone();
48                async move { discovery }
49            }),
50        );
51    }
52
53    let proxy_for_fallback = proxy.clone();
54    router
55        .route("/responses", responses_websocket_route(proxy.clone()))
56        .route("/v1/responses", responses_websocket_route(proxy.clone()))
57        .route(
58            "/backend-api/codex/responses",
59            responses_websocket_route(proxy),
60        )
61        .route(
62            "/{*path}",
63            any(move |req| handle_proxy(proxy_for_fallback.clone(), req)),
64        )
65        .layer(middleware::from_fn(reject_admin_paths_from_proxy))
66}
67
68fn responses_websocket_route(proxy: ProxyService) -> axum::routing::MethodRouter {
69    let proxy_for_ws = proxy.clone();
70    let proxy_for_fallback = proxy;
71    on(
72        MethodFilter::GET,
73        move |ws: Result<WebSocketUpgrade, WebSocketUpgradeRejection>,
74              headers: HeaderMap,
75              uri: Uri| {
76            handle_responses_websocket(proxy_for_ws.clone(), ws, headers, uri)
77        },
78    )
79    .fallback(move |req| handle_proxy(proxy_for_fallback.clone(), req))
80}
81
82pub fn admin_listener_router(proxy: ProxyService) -> Router {
83    router(proxy).layer(middleware::from_fn(require_admin_path_only))
84}
codex_helper_core/proxy/router_setup.rs

codex_helper_core/proxy/
router_setup.rs
