shift-proxy 0.10.2

Native Rust HTTP proxy for SHIFT — intercepts AI API requests, optimizes image payloads, and forwards to upstream providers
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

//! OpenAI route handler — POST /v1/chat/completions
//!
//! Intercepts OpenAI API requests, runs the SHIFT optimization pipeline
//! on the payload, records stats with full per-image token savings, and
//! forwards to the real OpenAI API.
//!
//! The CPU-intensive optimization runs on a blocking thread to avoid
//! starving the tokio event loop.

use crate::body::extract_body;
use crate::forward::forward_request;
use crate::optimize::optimize_payload;
use crate::ProxyState;
use axum::body::Bytes;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode, Uri};
use axum::response::{IntoResponse, Response};

/// POST /v1/chat/completions — optimize and forward to OpenAI.
pub async fn openai_handler(
    State(state): State<ProxyState>,
    uri: Uri,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    let body = match extract_body(&headers, body) {
        Ok(s) => s,
        Err(e) => {
            return (
                StatusCode::BAD_REQUEST,
                axum::Json(serde_json::json!({"error": e})),
            )
                .into_response();
        }
    };
    let config = state.config.shift_config("openai");
    let base_url = &state.config.providers.openai;
    let query = uri.query().map(|q| format!("?{}", q)).unwrap_or_default();
    let target_url = format!("{}{}{}", base_url, uri.path(), query);

    // Run SHIFT optimization pipeline on a blocking thread to avoid
    // starving the async runtime during CPU-intensive image operations.
    let start = std::time::Instant::now();
    let body_clone = body.clone();
    let optimization_result =
        tokio::task::spawn_blocking(move || optimize_payload(&body_clone, &config)).await;

    let (final_body, optimized) = match optimization_result {
        Ok(Some((transformed_json, report))) => {
            let duration_ms = start.elapsed().as_millis() as u64;

            // Record session stats (in-memory)
            state.session.record(&report);

            // Record persistent stats with FULL token savings
            let record = shift_preflight::stats::record_from_report(&report, "openai", duration_ms);
            if let Err(e) = shift_preflight::stats::record_run(&record, None) {
                tracing::warn!("failed to save stats: {}", e);
            }

            if state.config.verbose {
                let saved = report.original_size.saturating_sub(report.transformed_size);
                if saved > 0 {
                    tracing::info!(
                        "OpenAI: saved {:.1}KB ({} tokens)",
                        saved as f64 / 1024.0,
                        report.token_savings.openai_saved(),
                    );
                }
            }

            (transformed_json, true)
        }
        Ok(None) | Err(_) => (body, false),
    };

    if state.config.verbose && !optimized {
        tracing::debug!("OpenAI: no optimization applied (passthrough)");
    }

    forward_request(
        &state.http_client,
        "POST",
        &target_url,
        &headers,
        Some(final_body),
    )
    .await
}