shift_proxy/
state.rs

1//! Shared proxy state and configuration.
2
3use shift_preflight::{DriveMode, ShiftConfig, SvgMode};
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::{Arc, Mutex};
6use std::time::Instant;
7
8/// Proxy configuration.
9#[derive(Debug, Clone)]
10pub struct ProxyConfig {
11    pub port: u16,
12    pub mode: DriveMode,
13    pub verbose: bool,
14    /// Custom upstream provider URLs (override defaults).
15    pub providers: ProviderUrls,
16}
17
18#[derive(Debug, Clone)]
19pub struct ProviderUrls {
20    pub anthropic: String,
21    pub openai: String,
22    pub google: String,
23}
24
25impl Default for ProviderUrls {
26    fn default() -> Self {
27        Self {
28            anthropic: "https://api.anthropic.com".to_string(),
29            openai: "https://api.openai.com".to_string(),
30            google: "https://generativelanguage.googleapis.com".to_string(),
31        }
32    }
33}
34
35impl Default for ProxyConfig {
36    fn default() -> Self {
37        Self {
38            port: 8787,
39            mode: DriveMode::Balanced,
40            verbose: false,
41            providers: ProviderUrls::default(),
42        }
43    }
44}
45
46impl ProxyConfig {
47    /// Build a `ShiftConfig` for the optimization pipeline.
48    pub fn shift_config(&self, provider: &str) -> ShiftConfig {
49        ShiftConfig {
50            mode: self.mode,
51            svg_mode: SvgMode::Raster,
52            provider: provider.to_string(),
53            model: None,
54            dry_run: false,
55            verbose: self.verbose,
56            profile_path: None,
57            limits: shift_preflight::SafetyLimits::default(),
58        }
59    }
60}
61
62/// Shared proxy state, threaded through axum handlers via `State<ProxyState>`.
63#[derive(Clone)]
64pub struct ProxyState {
65    pub config: ProxyConfig,
66    pub http_client: reqwest::Client,
67    pub session: Arc<SessionStats>,
68}
69
70impl ProxyState {
71    pub fn new(config: ProxyConfig) -> Self {
72        let http_client = reqwest::Client::builder()
73            // Use connect_timeout, not total timeout — streaming responses
74            // (SSE from Anthropic/OpenAI) can run for minutes.
75            .connect_timeout(std::time::Duration::from_secs(30))
76            // SECURITY: Do not follow redirects. Prevents SSRF via redirect
77            // to internal services (e.g., cloud metadata endpoints).
78            .redirect(reqwest::redirect::Policy::none())
79            .build()
80            .expect("failed to build HTTP client");
81
82        Self {
83            config,
84            http_client,
85            session: Arc::new(SessionStats::new()),
86        }
87    }
88}
89
90/// In-memory session statistics (mirrors TS SessionStats).
91pub struct SessionStats {
92    pub started_at: Instant,
93    pub total_requests: AtomicU64,
94    pub total_images: AtomicU64,
95    pub total_images_modified: AtomicU64,
96    pub total_bytes_saved: AtomicU64,
97    pub token_savings: Mutex<TokenSavingsAccum>,
98}
99
100#[derive(Debug, Clone, Default, serde::Serialize)]
101pub struct TokenSavingsAccum {
102    pub openai_before: u64,
103    pub openai_after: u64,
104    pub anthropic_before: u64,
105    pub anthropic_after: u64,
106}
107
108impl Default for SessionStats {
109    fn default() -> Self {
110        Self::new()
111    }
112}
113
114impl SessionStats {
115    pub fn new() -> Self {
116        Self {
117            started_at: Instant::now(),
118            total_requests: AtomicU64::new(0),
119            total_images: AtomicU64::new(0),
120            total_images_modified: AtomicU64::new(0),
121            total_bytes_saved: AtomicU64::new(0),
122            token_savings: Mutex::new(TokenSavingsAccum::default()),
123        }
124    }
125
126    /// Record stats from a completed optimization run.
127    ///
128    /// Atomics use `Ordering::Relaxed` because these are independent counters
129    /// with no happens-before relationship — approximate consistency is fine
130    /// for diagnostic stats.
131    pub fn record(&self, report: &shift_preflight::Report) {
132        self.total_requests.fetch_add(1, Ordering::Relaxed);
133        self.total_images
134            .fetch_add(report.images_found as u64, Ordering::Relaxed);
135        self.total_images_modified
136            .fetch_add(report.images_modified as u64, Ordering::Relaxed);
137        let saved = report.original_size.saturating_sub(report.transformed_size) as u64;
138        self.total_bytes_saved.fetch_add(saved, Ordering::Relaxed);
139
140        // Recover from mutex poisoning — the inner data (simple counters)
141        // has no invariants that could be violated by a panic.
142        let mut ts = self.token_savings.lock().unwrap_or_else(|e| e.into_inner());
143        ts.openai_before += report.token_savings.openai_before;
144        ts.openai_after += report.token_savings.openai_after;
145        ts.anthropic_before += report.token_savings.anthropic_before;
146        ts.anthropic_after += report.token_savings.anthropic_after;
147    }
148
149    /// Serialize to JSON for the /stats endpoint.
150    pub fn to_json(&self) -> serde_json::Value {
151        let ts = self
152            .token_savings
153            .lock()
154            .unwrap_or_else(|e| e.into_inner())
155            .clone();
156        serde_json::json!({
157            "startedAt": format!("{:.0?}", self.started_at.elapsed()),
158            "totalRequests": self.total_requests.load(Ordering::Relaxed),
159            "totalImages": self.total_images.load(Ordering::Relaxed),
160            "totalImagesModified": self.total_images_modified.load(Ordering::Relaxed),
161            "totalBytesSaved": self.total_bytes_saved.load(Ordering::Relaxed),
162            "tokenSavings": {
163                "openai_before": ts.openai_before,
164                "openai_after": ts.openai_after,
165                "anthropic_before": ts.anthropic_before,
166                "anthropic_after": ts.anthropic_after,
167            }
168        })
169    }
170}
shift_proxy/state.rs

shift_proxy/
state.rs
