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_edits, 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(
71            "/images/edits",
72            on(MethodFilter::POST, {
73                let proxy = proxy.clone();
74                move |req| handle_openai_images_edits(proxy.clone(), req)
75            }),
76        )
77        .route(
78            "/v1/images/edits",
79            on(MethodFilter::POST, {
80                let proxy = proxy.clone();
81                move |req| handle_openai_images_edits(proxy.clone(), req)
82            }),
83        )
84        .route("/responses", responses_websocket_route(proxy.clone()))
85        .route("/v1/responses", responses_websocket_route(proxy.clone()))
86        .route(
87            "/backend-api/codex/responses",
88            responses_websocket_route(proxy),
89        )
90        .route(
91            "/{*path}",
92            any(move |req| handle_proxy(proxy_for_fallback.clone(), req)),
93        )
94        .layer(middleware::from_fn(reject_admin_paths_from_proxy))
95}
96
97fn responses_websocket_route(proxy: ProxyService) -> axum::routing::MethodRouter {
98    let proxy_for_ws = proxy.clone();
99    let proxy_for_fallback = proxy;
100    on(
101        MethodFilter::GET,
102        move |ws: Result<WebSocketUpgrade, WebSocketUpgradeRejection>,
103              headers: HeaderMap,
104              uri: Uri| {
105            handle_responses_websocket(proxy_for_ws.clone(), ws, headers, uri)
106        },
107    )
108    .fallback(move |req| handle_proxy(proxy_for_fallback.clone(), req))
109}
110
111pub fn admin_listener_router(proxy: ProxyService) -> Router {
112    router(proxy).layer(middleware::from_fn(require_admin_path_only))
113}
codex_helper_core/proxy/router_setup.rs

codex_helper_core/proxy/
router_setup.rs
