codex_helper_core/
runtime_host.rs

1use std::collections::HashMap;
2use std::net::{IpAddr, SocketAddr};
3use std::sync::{Arc, Mutex};
4
5use anyhow::Result;
6use axum::Router;
7use reqwest::Client;
8use tokio::sync::watch;
9use tokio::task::JoinHandle;
10
11use crate::config::{
12    LoadedProxyConfig, ProxyConfig, ServiceKind, load_or_bootstrap_for_service_with_v4_source,
13    model_routing_warnings,
14};
15use crate::lb::LbState;
16use crate::proxy::{
17    ProxyService, admin_listener_router, admin_loopback_addr_for_proxy_port, local_proxy_base_url,
18    proxy_only_router_with_admin_base_url,
19};
20use crate::state::ProxyState;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum ProxyLifetimeMode {
24    Ephemeral,
25    Resident,
26}
27
28impl ProxyLifetimeMode {
29    pub fn is_resident(self) -> bool {
30        matches!(self, Self::Resident)
31    }
32}
33
34pub struct ProxyRuntime {
35    pub service_name: &'static str,
36    pub host: IpAddr,
37    pub port: u16,
38    pub admin_addr: SocketAddr,
39    pub config: Arc<ProxyConfig>,
40    pub proxy: ProxyService,
41    pub state: Arc<ProxyState>,
42    pub shutdown_tx: watch::Sender<bool>,
43    shutdown_rx: watch::Receiver<bool>,
44    listener: Option<tokio::net::TcpListener>,
45    admin_listener: Option<tokio::net::TcpListener>,
46    app: Option<Router>,
47    admin_app: Option<Router>,
48}
49
50pub struct RunningProxyRuntime {
51    pub service_name: &'static str,
52    pub host: IpAddr,
53    pub port: u16,
54    pub admin_addr: SocketAddr,
55    pub config: Arc<ProxyConfig>,
56    pub proxy: ProxyService,
57    pub state: Arc<ProxyState>,
58    pub shutdown_tx: watch::Sender<bool>,
59    pub server_handle: JoinHandle<Result<()>>,
60}
61
62impl ProxyRuntime {
63    pub fn shutdown_receiver(&self) -> watch::Receiver<bool> {
64        self.shutdown_rx.clone()
65    }
66
67    pub fn start(mut self) -> RunningProxyRuntime {
68        let listener = self
69            .listener
70            .take()
71            .expect("proxy runtime listener should exist before start");
72        let admin_listener = self
73            .admin_listener
74            .take()
75            .expect("proxy runtime admin listener should exist before start");
76        let app = self
77            .app
78            .take()
79            .expect("proxy runtime app should exist before start");
80        let admin_app = self
81            .admin_app
82            .take()
83            .expect("proxy runtime admin app should exist before start");
84        let shutdown_rx = self.shutdown_rx.clone();
85        let server_handle =
86            spawn_proxy_runtime_servers(listener, admin_listener, app, admin_app, shutdown_rx);
87
88        RunningProxyRuntime {
89            service_name: self.service_name,
90            host: self.host,
91            port: self.port,
92            admin_addr: self.admin_addr,
93            config: self.config,
94            proxy: self.proxy,
95            state: self.state,
96            shutdown_tx: self.shutdown_tx,
97            server_handle,
98        }
99    }
100}
101
102impl RunningProxyRuntime {
103    pub async fn wait(self) -> Result<()> {
104        self.server_handle
105            .await
106            .map_err(|error| anyhow::anyhow!("server task join error: {error}"))?
107    }
108}
109
110pub async fn build_proxy_runtime(
111    service_kind: ServiceKind,
112    host: IpAddr,
113    port: u16,
114) -> Result<ProxyRuntime> {
115    let service_name = service_name_for_kind(service_kind);
116    let loaded = load_or_bootstrap_for_service_with_v4_source(service_kind).await?;
117    build_proxy_runtime_from_loaded(service_name, host, port, loaded).await
118}
119
120pub async fn build_proxy_runtime_from_loaded(
121    service_name: &'static str,
122    host: IpAddr,
123    port: u16,
124    loaded: LoadedProxyConfig,
125) -> Result<ProxyRuntime> {
126    let addr: SocketAddr = SocketAddr::from((host, port));
127    let admin_addr = admin_loopback_addr_for_proxy_port(port);
128    let listener = tokio::net::TcpListener::bind(addr).await?;
129    let admin_listener = tokio::net::TcpListener::bind(admin_addr).await?;
130    build_proxy_runtime_from_bound_listeners(
131        service_name,
132        host,
133        port,
134        loaded,
135        listener,
136        admin_listener,
137    )
138    .await
139}
140
141pub async fn build_proxy_runtime_from_bound_listeners(
142    service_name: &'static str,
143    host: IpAddr,
144    port: u16,
145    loaded: LoadedProxyConfig,
146    listener: tokio::net::TcpListener,
147    admin_listener: tokio::net::TcpListener,
148) -> Result<ProxyRuntime> {
149    let cfg = loaded.runtime;
150    validate_service_has_upstream(service_name, &cfg)?;
151    let v4_source = loaded.v4.map(Arc::new);
152
153    let warnings = model_routing_warnings(&cfg, service_name);
154    if !warnings.is_empty() {
155        tracing::warn!("======== Model routing config warnings ========");
156        for warning in warnings {
157            tracing::warn!("{}", warning);
158        }
159        tracing::warn!("==============================================");
160    }
161
162    let client = Client::builder()
163        .connect_timeout(std::time::Duration::from_secs(10))
164        .tcp_keepalive(std::time::Duration::from_secs(30))
165        .pool_idle_timeout(std::time::Duration::from_secs(30))
166        .build()?;
167
168    let lb_states = Arc::new(Mutex::new(HashMap::<String, LbState>::new()));
169    let admin_addr = admin_loopback_addr_for_proxy_port(port);
170    let cfg = Arc::new(cfg);
171    let (shutdown_tx, shutdown_rx) = watch::channel(false);
172
173    let proxy = ProxyService::new_with_v4_source_and_shutdown(
174        client,
175        cfg.clone(),
176        v4_source,
177        service_name,
178        lb_states,
179        Some(shutdown_tx.clone()),
180    );
181    let state = proxy.state_handle();
182    let app = proxy_only_router_with_admin_base_url(
183        proxy.clone(),
184        Some(local_proxy_base_url(admin_addr.port())),
185    );
186    let admin_app = admin_listener_router(proxy.clone());
187
188    Ok(ProxyRuntime {
189        service_name,
190        host,
191        port,
192        admin_addr,
193        config: cfg,
194        proxy,
195        state,
196        shutdown_tx,
197        shutdown_rx,
198        listener: Some(listener),
199        admin_listener: Some(admin_listener),
200        app: Some(app),
201        admin_app: Some(admin_app),
202    })
203}
204
205pub fn service_name_for_kind(service_kind: ServiceKind) -> &'static str {
206    match service_kind {
207        ServiceKind::Codex => "codex",
208        ServiceKind::Claude => "claude",
209    }
210}
211
212pub fn service_kind_for_name(service_name: &str) -> ServiceKind {
213    match service_name {
214        "claude" => ServiceKind::Claude,
215        _ => ServiceKind::Codex,
216    }
217}
218
219pub fn validate_service_has_upstream(service_name: &str, cfg: &ProxyConfig) -> Result<()> {
220    match service_name {
221        "claude" => {
222            if !cfg.claude.has_stations() || cfg.claude.active_station().is_none() {
223                anyhow::bail!(
224                    "未找到任何可用的 Claude 上游配置，请先确保 ~/.claude/settings.json 配置完整，\
225或在 ~/.codex-helper/config.toml（或 config.json）的 `claude` 段下手动添加上游配置"
226                );
227            }
228        }
229        _ => {
230            if !cfg.codex.has_stations() || cfg.codex.active_station().is_none() {
231                anyhow::bail!(
232                    "未找到任何可用的 Codex 上游配置，请先确保 ~/.codex/config.toml 与 ~/.codex/auth.json 配置完整，或手动编辑 ~/.codex-helper/config.toml（或 config.json）添加配置"
233                );
234            }
235        }
236    }
237    Ok(())
238}
239
240fn spawn_proxy_runtime_servers(
241    listener: tokio::net::TcpListener,
242    admin_listener: tokio::net::TcpListener,
243    app: Router,
244    admin_app: Router,
245    shutdown_rx: watch::Receiver<bool>,
246) -> JoinHandle<Result<()>> {
247    let proxy_server_shutdown = {
248        let mut rx = shutdown_rx.clone();
249        async move {
250            let _ = rx.changed().await;
251        }
252    };
253    let admin_server_shutdown = {
254        let mut rx = shutdown_rx;
255        async move {
256            let _ = rx.changed().await;
257        }
258    };
259
260    tokio::spawn(async move {
261        tokio::try_join!(
262            axum::serve(
263                listener,
264                app.into_make_service_with_connect_info::<SocketAddr>(),
265            )
266            .with_graceful_shutdown(proxy_server_shutdown),
267            axum::serve(
268                admin_listener,
269                admin_app.into_make_service_with_connect_info::<SocketAddr>(),
270            )
271            .with_graceful_shutdown(admin_server_shutdown),
272        )?;
273        Ok(())
274    })
275}
codex_helper_core/runtime_host.rs

codex_helper_core/
runtime_host.rs
