Skip to main content

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.
37pub async fn start_server(config: ProxyConfig) -> anyhow::Result<()> {
38    // Initialize tracing subscriber so that tracing::warn!/error!/info!
39    // calls in route handlers are actually visible on stderr.
40    let filter = if config.verbose {
41        "shift_proxy=debug,tower_http=debug"
42    } else {
43        "shift_proxy=warn"
44    };
45    tracing_subscriber::fmt()
46        .with_env_filter(
47            tracing_subscriber::EnvFilter::try_from_default_env()
48                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(filter)),
49        )
50        .with_target(false)
51        .init();
52
53    let port = config.port;
54    let verbose = config.verbose;
55    let app = create_app(config);
56
57    let addr = SocketAddr::from(([127, 0, 0, 1], port));
58    let listener = TcpListener::bind(addr).await?;
59
60    if verbose {
61        tracing::info!("shift proxy listening on http://{}", addr);
62    }
63    eprintln!("[shift] proxy listening on http://{}", addr);
64
65    axum::serve(listener, app)
66        .with_graceful_shutdown(shutdown_signal())
67        .await?;
68
69    Ok(())
70}
71
72async fn shutdown_signal() {
73    let ctrl_c = async {
74        tokio::signal::ctrl_c()
75            .await
76            .expect("failed to install Ctrl+C handler");
77    };
78
79    #[cfg(unix)]
80    let terminate = async {
81        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
82            .expect("failed to install SIGTERM handler")
83            .recv()
84            .await;
85    };
86
87    #[cfg(not(unix))]
88    let terminate = std::future::pending::<()>();
89
90    tokio::select! {
91        _ = ctrl_c => {},
92        _ = terminate => {},
93    }
94}