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::host_local::HostLocalSessionHistoryMode;
16use crate::lb::LbState;
17use crate::proxy::{
18 ProxyService, admin_listener_router, admin_loopback_addr_for_proxy_port,
19 proxy_only_router_with_admin_base_url,
20};
21use crate::state::ProxyState;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum ProxyLifetimeMode {
25 Ephemeral,
26 Resident,
27}
28
29impl ProxyLifetimeMode {
30 pub fn is_resident(self) -> bool {
31 matches!(self, Self::Resident)
32 }
33}
34
35pub struct ProxyRuntime {
36 pub service_name: &'static str,
37 pub host: IpAddr,
38 pub port: u16,
39 pub admin_addr: SocketAddr,
40 pub config: Arc<ProxyConfig>,
41 pub proxy: ProxyService,
42 pub state: Arc<ProxyState>,
43 pub shutdown_tx: watch::Sender<bool>,
44 shutdown_rx: watch::Receiver<bool>,
45 listener: Option<tokio::net::TcpListener>,
46 admin_listener: Option<tokio::net::TcpListener>,
47 app: Option<Router>,
48 admin_app: Option<Router>,
49}
50
51pub struct RunningProxyRuntime {
52 pub service_name: &'static str,
53 pub host: IpAddr,
54 pub port: u16,
55 pub admin_addr: SocketAddr,
56 pub config: Arc<ProxyConfig>,
57 pub proxy: ProxyService,
58 pub state: Arc<ProxyState>,
59 pub shutdown_tx: watch::Sender<bool>,
60 pub server_handle: JoinHandle<Result<()>>,
61}
62
63#[derive(Debug, Clone)]
64pub struct ProxyRuntimeOptions {
65 pub admin_addr: SocketAddr,
66 pub advertised_admin_base_url: Option<String>,
67 pub host_local_session_history_mode: HostLocalSessionHistoryMode,
68}
69
70impl ProxyRuntimeOptions {
71 pub fn for_proxy_port(port: u16) -> Self {
72 let admin_addr = admin_loopback_addr_for_proxy_port(port);
73 Self {
74 admin_addr,
75 advertised_admin_base_url: admin_discovery_base_url(admin_addr, None),
76 host_local_session_history_mode: HostLocalSessionHistoryMode::Auto,
77 }
78 }
79
80 pub fn with_admin_addr(mut self, admin_addr: SocketAddr) -> Self {
81 self.admin_addr = admin_addr;
82 self.advertised_admin_base_url = admin_discovery_base_url(admin_addr, None);
83 self
84 }
85
86 pub fn with_advertised_admin_base_url(
87 mut self,
88 advertised_admin_base_url: Option<String>,
89 ) -> Self {
90 self.advertised_admin_base_url =
91 admin_discovery_base_url(self.admin_addr, advertised_admin_base_url.as_deref());
92 self
93 }
94
95 pub fn with_host_local_session_history_mode(
96 mut self,
97 mode: HostLocalSessionHistoryMode,
98 ) -> Self {
99 self.host_local_session_history_mode = mode;
100 self
101 }
102}
103
104impl ProxyRuntime {
105 pub fn shutdown_receiver(&self) -> watch::Receiver<bool> {
106 self.shutdown_rx.clone()
107 }
108
109 pub fn start(mut self) -> RunningProxyRuntime {
110 let listener = self
111 .listener
112 .take()
113 .expect("proxy runtime listener should exist before start");
114 let admin_listener = self
115 .admin_listener
116 .take()
117 .expect("proxy runtime admin listener should exist before start");
118 let app = self
119 .app
120 .take()
121 .expect("proxy runtime app should exist before start");
122 let admin_app = self
123 .admin_app
124 .take()
125 .expect("proxy runtime admin app should exist before start");
126 let shutdown_rx = self.shutdown_rx.clone();
127 let server_handle =
128 spawn_proxy_runtime_servers(listener, admin_listener, app, admin_app, shutdown_rx);
129
130 RunningProxyRuntime {
131 service_name: self.service_name,
132 host: self.host,
133 port: self.port,
134 admin_addr: self.admin_addr,
135 config: self.config,
136 proxy: self.proxy,
137 state: self.state,
138 shutdown_tx: self.shutdown_tx,
139 server_handle,
140 }
141 }
142}
143
144impl RunningProxyRuntime {
145 pub async fn wait(self) -> Result<()> {
146 self.server_handle
147 .await
148 .map_err(|error| anyhow::anyhow!("server task join error: {error}"))?
149 }
150}
151
152pub async fn build_proxy_runtime(
153 service_kind: ServiceKind,
154 host: IpAddr,
155 port: u16,
156) -> Result<ProxyRuntime> {
157 let service_name = service_name_for_kind(service_kind);
158 let loaded = load_or_bootstrap_for_service_with_v4_source(service_kind).await?;
159 build_proxy_runtime_from_loaded(service_name, host, port, loaded).await
160}
161
162pub async fn build_proxy_runtime_from_loaded(
163 service_name: &'static str,
164 host: IpAddr,
165 port: u16,
166 loaded: LoadedProxyConfig,
167) -> Result<ProxyRuntime> {
168 build_proxy_runtime_from_loaded_with_options(
169 service_name,
170 host,
171 port,
172 ProxyRuntimeOptions::for_proxy_port(port),
173 loaded,
174 )
175 .await
176}
177
178pub async fn build_proxy_runtime_from_loaded_with_admin_addr(
179 service_name: &'static str,
180 host: IpAddr,
181 port: u16,
182 admin_addr: SocketAddr,
183 loaded: LoadedProxyConfig,
184) -> Result<ProxyRuntime> {
185 build_proxy_runtime_from_loaded_with_options(
186 service_name,
187 host,
188 port,
189 ProxyRuntimeOptions::for_proxy_port(port).with_admin_addr(admin_addr),
190 loaded,
191 )
192 .await
193}
194
195pub async fn build_proxy_runtime_from_loaded_with_options(
196 service_name: &'static str,
197 host: IpAddr,
198 port: u16,
199 options: ProxyRuntimeOptions,
200 loaded: LoadedProxyConfig,
201) -> Result<ProxyRuntime> {
202 let addr: SocketAddr = SocketAddr::from((host, port));
203 let listener = tokio::net::TcpListener::bind(addr).await?;
204 let admin_listener = tokio::net::TcpListener::bind(options.admin_addr).await?;
205 build_proxy_runtime_from_bound_listeners_with_options(
206 service_name,
207 host,
208 port,
209 options,
210 loaded,
211 listener,
212 admin_listener,
213 )
214 .await
215}
216
217pub async fn build_proxy_runtime_from_bound_listeners(
218 service_name: &'static str,
219 host: IpAddr,
220 port: u16,
221 loaded: LoadedProxyConfig,
222 listener: tokio::net::TcpListener,
223 admin_listener: tokio::net::TcpListener,
224) -> Result<ProxyRuntime> {
225 build_proxy_runtime_from_bound_listeners_with_options(
226 service_name,
227 host,
228 port,
229 ProxyRuntimeOptions::for_proxy_port(port),
230 loaded,
231 listener,
232 admin_listener,
233 )
234 .await
235}
236
237pub async fn build_proxy_runtime_from_bound_listeners_with_admin_addr(
238 service_name: &'static str,
239 host: IpAddr,
240 port: u16,
241 admin_addr: SocketAddr,
242 loaded: LoadedProxyConfig,
243 listener: tokio::net::TcpListener,
244 admin_listener: tokio::net::TcpListener,
245) -> Result<ProxyRuntime> {
246 build_proxy_runtime_from_bound_listeners_with_options(
247 service_name,
248 host,
249 port,
250 ProxyRuntimeOptions::for_proxy_port(port).with_admin_addr(admin_addr),
251 loaded,
252 listener,
253 admin_listener,
254 )
255 .await
256}
257
258pub async fn build_proxy_runtime_from_bound_listeners_with_options(
259 service_name: &'static str,
260 host: IpAddr,
261 port: u16,
262 options: ProxyRuntimeOptions,
263 loaded: LoadedProxyConfig,
264 listener: tokio::net::TcpListener,
265 admin_listener: tokio::net::TcpListener,
266) -> Result<ProxyRuntime> {
267 let cfg = loaded.runtime;
268 validate_service_has_upstream(service_name, &cfg)?;
269 let v4_source = loaded.v4.map(Arc::new);
270
271 let warnings = model_routing_warnings(&cfg, service_name);
272 if !warnings.is_empty() {
273 tracing::warn!("======== Model routing config warnings ========");
274 for warning in warnings {
275 tracing::warn!("{}", warning);
276 }
277 tracing::warn!("==============================================");
278 }
279
280 let client = Client::builder()
281 .connect_timeout(std::time::Duration::from_secs(10))
282 .tcp_keepalive(std::time::Duration::from_secs(30))
283 .pool_idle_timeout(std::time::Duration::from_secs(30))
284 .build()?;
285
286 let lb_states = Arc::new(Mutex::new(HashMap::<String, LbState>::new()));
287 let cfg = Arc::new(cfg);
288 let (shutdown_tx, shutdown_rx) = watch::channel(false);
289
290 let proxy = ProxyService::new_with_v4_source_and_shutdown(
291 client,
292 cfg.clone(),
293 v4_source,
294 service_name,
295 lb_states,
296 Some(shutdown_tx.clone()),
297 )
298 .with_host_local_session_history_mode(options.host_local_session_history_mode);
299 let state = proxy.state_handle();
300 let app = proxy_only_router_with_admin_base_url(
301 proxy.clone(),
302 options.advertised_admin_base_url.clone(),
303 );
304 let admin_app = admin_listener_router(proxy.clone());
305
306 Ok(ProxyRuntime {
307 service_name,
308 host,
309 port,
310 admin_addr: options.admin_addr,
311 config: cfg,
312 proxy,
313 state,
314 shutdown_tx,
315 shutdown_rx,
316 listener: Some(listener),
317 admin_listener: Some(admin_listener),
318 app: Some(app),
319 admin_app: Some(admin_app),
320 })
321}
322
323fn admin_discovery_base_url(
324 admin_addr: SocketAddr,
325 advertised_admin_base_url: Option<&str>,
326) -> Option<String> {
327 if let Some(url) = advertised_admin_base_url
328 .map(str::trim)
329 .filter(|url| !url.is_empty())
330 {
331 return Some(url.trim_end_matches('/').to_string());
332 }
333 if admin_addr.ip().is_unspecified() {
334 None
335 } else {
336 Some(format!("http://{admin_addr}"))
337 }
338}
339
340pub fn service_name_for_kind(service_kind: ServiceKind) -> &'static str {
341 match service_kind {
342 ServiceKind::Codex => "codex",
343 ServiceKind::Claude => "claude",
344 }
345}
346
347pub fn service_kind_for_name(service_name: &str) -> ServiceKind {
348 match service_name {
349 "claude" => ServiceKind::Claude,
350 _ => ServiceKind::Codex,
351 }
352}
353
354pub fn validate_service_has_upstream(service_name: &str, cfg: &ProxyConfig) -> Result<()> {
355 match service_name {
356 "claude" => {
357 if !cfg.claude.has_stations() || cfg.claude.active_station().is_none() {
358 anyhow::bail!(
359 "未找到任何可用的 Claude 上游配置，请先确保 ~/.claude/settings.json 配置完整，\
360或在 ~/.codex-helper/config.toml（或 config.json）的 `claude` 段下手动添加上游配置"
361 );
362 }
363 }
364 _ => {
365 if !cfg.codex.has_stations() || cfg.codex.active_station().is_none() {
366 anyhow::bail!(
367 "未找到任何可用的 Codex 上游配置，请先确保 ~/.codex/config.toml 与 ~/.codex/auth.json 配置完整，或手动编辑 ~/.codex-helper/config.toml（或 config.json）添加配置"
368 );
369 }
370 }
371 }
372 Ok(())
373}
374
375fn spawn_proxy_runtime_servers(
376 listener: tokio::net::TcpListener,
377 admin_listener: tokio::net::TcpListener,
378 app: Router,
379 admin_app: Router,
380 shutdown_rx: watch::Receiver<bool>,
381) -> JoinHandle<Result<()>> {
382 let proxy_server_shutdown = {
383 let mut rx = shutdown_rx.clone();
384 async move {
385 let _ = rx.changed().await;
386 }
387 };
388 let admin_server_shutdown = {
389 let mut rx = shutdown_rx;
390 async move {
391 let _ = rx.changed().await;
392 }
393 };
394
395 tokio::spawn(async move {
396 tokio::try_join!(
397 axum::serve(
398 listener,
399 app.into_make_service_with_connect_info::<SocketAddr>(),
400 )
401 .with_graceful_shutdown(proxy_server_shutdown),
402 axum::serve(
403 admin_listener,
404 admin_app.into_make_service_with_connect_info::<SocketAddr>(),
405 )
406 .with_graceful_shutdown(admin_server_shutdown),
407 )?;
408 Ok(())
409 })
410}
411
412#[cfg(test)]
413mod tests {
414 use super::*;
415
416 #[test]
417 fn admin_discovery_url_is_not_advertised_for_unspecified_bind() {
418 let admin_addr = SocketAddr::from(([0, 0, 0, 0], 4211));
419
420 assert_eq!(admin_discovery_base_url(admin_addr, None), None);
421 }
422
423 #[test]
424 fn admin_discovery_url_uses_explicit_bind_address() {
425 let admin_addr = SocketAddr::from(([192, 168, 1, 10], 4211));
426
427 assert_eq!(
428 admin_discovery_base_url(admin_addr, None),
429 Some("http://192.168.1.10:4211".to_string())
430 );
431 }
432
433 #[test]
434 fn admin_discovery_url_uses_advertised_url_when_provided() {
435 let admin_addr = SocketAddr::from(([0, 0, 0, 0], 4211));
436
437 assert_eq!(
438 admin_discovery_base_url(admin_addr, Some("http://nas.local:4211/")),
439 Some("http://nas.local:4211".to_string())
440 );
441 }
442}