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

codex_helper_core/proxy/
router_setup.rs
