shift_proxy/
lib.rs

1//! SHIFT native proxy — Rust HTTP server that intercepts AI API requests,
2//! optimizes image payloads via `shift-preflight`, and forwards to upstream
3//! providers. Replaces the Node.js/Hono proxy with a single-binary server.
4//!
5//! ## Architecture
6//!
7//! ```text
8//! Client (OpenCode, Claude Code, Codex, etc.)
9//!   │
10//!   ├── POST /v1/messages         → Anthropic (optimize + forward)
11//!   ├── POST /messages            → Anthropic (rewrite → /v1/messages)
12//!   ├── POST /v1/chat/completions → OpenAI   (optimize + forward)
13//!   ├── POST /v1beta/models/*     → Google   (passthrough)
14//!   ├── GET  /health              → Status
15//!   ├── GET  /stats               → Session stats
16//!   └── POST /*                   → Auto-detect provider (passthrough)
17//! ```
18
19pub mod forward;
20pub mod optimize;
21pub mod routes;
22pub mod state;
23
24use axum::Router;
25use std::net::SocketAddr;
26use tokio::net::TcpListener;
27
28pub use state::{ProxyConfig, ProxyState};
29
30/// Build the axum router with all proxy routes.
31pub fn create_app(config: ProxyConfig) -> Router {
32    let state = ProxyState::new(config);
33    routes::build_router(state)
34}
35
36/// Start the proxy server, blocking until shutdown signal.
37///
38/// Uses `axum::serve` which auto-negotiates HTTP/1.1 and HTTP/2 (h2c)
39/// via `hyper_util::server::conn::auto::Builder` internally, and provides
40/// graceful shutdown that drains in-flight connections.
41pub async fn start_server(config: ProxyConfig) -> anyhow::Result<()> {
42    // Initialize tracing subscriber so that tracing::warn!/error!/info!
43    // calls in route handlers are actually visible on stderr.
44    let filter = if config.verbose {
45        "shift_proxy=debug,tower_http=debug"
46    } else {
47        "shift_proxy=warn"
48    };
49    tracing_subscriber::fmt()
50        .with_env_filter(
51            tracing_subscriber::EnvFilter::try_from_default_env()
52                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(filter)),
53        )
54        .with_target(false)
55        .init();
56
57    let port = config.port;
58    let verbose = config.verbose;
59    let app = create_app(config);
60
61    let addr = SocketAddr::from(([127, 0, 0, 1], port));
62    let listener = TcpListener::bind(addr).await?;
63
64    if verbose {
65        tracing::info!("shift proxy listening on http://{}", addr);
66    }
67    eprintln!("[shift] proxy listening on http://{}", addr);
68
69    axum::serve(listener, app)
70        .with_graceful_shutdown(shutdown_signal())
71        .await?;
72
73    Ok(())
74}
75
76async fn shutdown_signal() {
77    let ctrl_c = async {
78        tokio::signal::ctrl_c()
79            .await
80            .expect("failed to install Ctrl+C handler");
81    };
82
83    #[cfg(unix)]
84    let terminate = async {
85        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
86            .expect("failed to install SIGTERM handler")
87            .recv()
88            .await;
89    };
90
91    #[cfg(not(unix))]
92    let terminate = std::future::pending::<()>();
93
94    tokio::select! {
95        _ = ctrl_c => {},
96        _ = terminate => {},
97    }
98}
shift_proxy/lib.rs

shift_proxy/
lib.rs
