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 /v1/chat/completions → OpenAI   (optimize + forward)
12//!   ├── POST /v1beta/models/*     → Google   (passthrough)
13//!   ├── GET  /health              → Status
14//!   ├── GET  /stats               → Session stats
15//!   └── POST /*                   → Auto-detect provider (passthrough)
16//! ```
17
18pub mod forward;
19pub mod optimize;
20pub mod routes;
21pub mod state;
22
23use axum::Router;
24use std::net::SocketAddr;
25use tokio::net::TcpListener;
26
27pub use state::{ProxyConfig, ProxyState};
28
29/// Build the axum router with all proxy routes.
30pub fn create_app(config: ProxyConfig) -> Router {
31    let state = ProxyState::new(config);
32    routes::build_router(state)
33}
34
35/// Start the proxy server, blocking until shutdown signal.
36pub async fn start_server(config: ProxyConfig) -> anyhow::Result<()> {
37    // Initialize tracing subscriber so that tracing::warn!/error!/info!
38    // calls in route handlers are actually visible on stderr.
39    let filter = if config.verbose {
40        "shift_proxy=debug,tower_http=debug"
41    } else {
42        "shift_proxy=warn"
43    };
44    tracing_subscriber::fmt()
45        .with_env_filter(
46            tracing_subscriber::EnvFilter::try_from_default_env()
47                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(filter)),
48        )
49        .with_target(false)
50        .init();
51
52    let port = config.port;
53    let verbose = config.verbose;
54    let app = create_app(config);
55
56    let addr = SocketAddr::from(([127, 0, 0, 1], port));
57    let listener = TcpListener::bind(addr).await?;
58
59    if verbose {
60        tracing::info!("shift proxy listening on http://{}", addr);
61    }
62    eprintln!("[shift] proxy listening on http://{}", addr);
63
64    axum::serve(listener, app)
65        .with_graceful_shutdown(shutdown_signal())
66        .await?;
67
68    Ok(())
69}
70
71async fn shutdown_signal() {
72    let ctrl_c = async {
73        tokio::signal::ctrl_c()
74            .await
75            .expect("failed to install Ctrl+C handler");
76    };
77
78    #[cfg(unix)]
79    let terminate = async {
80        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
81            .expect("failed to install SIGTERM handler")
82            .recv()
83            .await;
84    };
85
86    #[cfg(not(unix))]
87    let terminate = std::future::pending::<()>();
88
89    tokio::select! {
90        _ = ctrl_c => {},
91        _ = terminate => {},
92    }
93}