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 body;
20pub mod forward;
21pub mod optimize;
22pub mod routes;
23pub mod state;
24
25use axum::Router;
26use std::net::SocketAddr;
27use tokio::net::TcpListener;
28
29pub use state::{ProxyConfig, ProxyState};
30
31/// Build the axum router with all proxy routes.
32pub fn create_app(config: ProxyConfig) -> Router {
33    let state = ProxyState::new(config);
34    routes::build_router(state)
35}
36
37/// Start the proxy server, blocking until shutdown signal.
38///
39/// Uses `axum::serve` which auto-negotiates HTTP/1.1 and HTTP/2 (h2c)
40/// via `hyper_util::server::conn::auto::Builder` internally, and provides
41/// graceful shutdown that drains in-flight connections.
42pub async fn start_server(config: ProxyConfig) -> anyhow::Result<()> {
43    // Initialize tracing subscriber so that tracing::warn!/error!/info!
44    // calls in route handlers are actually visible on stderr.
45    let filter = if config.verbose {
46        "shift_proxy=debug,tower_http=debug"
47    } else {
48        "shift_proxy=warn"
49    };
50    tracing_subscriber::fmt()
51        .with_env_filter(
52            tracing_subscriber::EnvFilter::try_from_default_env()
53                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(filter)),
54        )
55        .with_target(false)
56        .init();
57
58    let port = config.port;
59    let verbose = config.verbose;
60    let app = create_app(config);
61
62    let addr = SocketAddr::from(([127, 0, 0, 1], port));
63    let listener = TcpListener::bind(addr).await?;
64
65    if verbose {
66        tracing::info!("shift proxy listening on http://{}", addr);
67    }
68    eprintln!("[shift] proxy listening on http://{}", addr);
69
70    axum::serve(listener, app)
71        .with_graceful_shutdown(shutdown_signal())
72        .await?;
73
74    Ok(())
75}
76
77async fn shutdown_signal() {
78    let ctrl_c = async {
79        tokio::signal::ctrl_c()
80            .await
81            .expect("failed to install Ctrl+C handler");
82    };
83
84    #[cfg(unix)]
85    let terminate = async {
86        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
87            .expect("failed to install SIGTERM handler")
88            .recv()
89            .await;
90    };
91
92    #[cfg(not(unix))]
93    let terminate = std::future::pending::<()>();
94
95    tokio::select! {
96        _ = ctrl_c => {},
97        _ = terminate => {},
98    }
99}
shift_proxy/lib.rs

shift_proxy/
lib.rs
