vibekit_proxy/
server.rs

1use std::collections::HashMap;
2use std::convert::Infallible;
3use std::net::SocketAddr;
4use std::sync::Arc;
5
6use hyper::service::{make_service_fn, service_fn};
7use hyper::{Body, Client, Method, Request, Response, Server, StatusCode};
8use hyper_tls::HttpsConnector;
9use serde_json::Value;
10use tokio::sync::{Mutex, RwLock};
11use tracing::{info, warn};
12
13use crate::config::ConfigManager;
14use crate::redaction::RedactionEngine;
15use crate::{Error, Result};
16
17#[derive(Debug, Clone)]
18struct SseAccumulator {
19    buffer: String,
20}
21
22pub struct ProxyServer {
23    port: u16,
24    request_count: Arc<Mutex<u64>>,
25    redaction_engine: RedactionEngine,
26    config_manager: Arc<RwLock<ConfigManager>>,
27    client: Client<HttpsConnector<hyper::client::HttpConnector>>,
28    sse_content_accumulators: Arc<RwLock<HashMap<u64, SseAccumulator>>>,
29}
30
31impl ProxyServer {
32    pub async fn new(port: u16, config_path: Option<String>) -> Result<Self> {
33        let https = HttpsConnector::new();
34        let client = Client::builder().build::<_, hyper::Body>(https);
35        
36        let mut config_manager = ConfigManager::new_with_path(config_path);
37        if let Err(e) = config_manager.load_config().await {
38            warn!("Warning: Could not load config file, using defaults: {}", e);
39        }
40
41        Ok(Self {
42            port,
43            request_count: Arc::new(Mutex::new(0)),
44            redaction_engine: RedactionEngine::new(),
45            config_manager: Arc::new(RwLock::new(config_manager)),
46            client,
47            sse_content_accumulators: Arc::new(RwLock::new(HashMap::new())),
48        })
49    }
50
51    pub async fn start(&self) -> Result<()> {
52        let addr = SocketAddr::from(([0, 0, 0, 0], self.port));
53        
54        let server_clone = Arc::new(self.clone());
55        let make_svc = make_service_fn(move |_conn| {
56            let server = Arc::clone(&server_clone);
57            async move {
58                Ok::<_, Infallible>(service_fn(move |req| {
59                    let server = Arc::clone(&server);
60                    async move { server.handle_request(req).await }
61                }))
62            }
63        });
64
65        let server = Server::bind(&addr).serve(make_svc);
66        info!("Proxy server listening on http://{}", addr);
67
68        if let Err(e) = server.await {
69            return Err(Error::Server(e.to_string()));
70        }
71
72        Ok(())
73    }
74
75    pub async fn stop(&self) {
76        info!("Stopping proxy server...");
77    }
78
79    async fn handle_request(&self, req: Request<Body>) -> std::result::Result<Response<Body>, Infallible> {
80        let result = self.handle_request_inner(req).await;
81        Ok(result.unwrap_or_else(|e| {
82            Response::builder()
83                .status(StatusCode::INTERNAL_SERVER_ERROR)
84                .body(Body::from(format!("Internal Server Error: {}", e)))
85                .unwrap()
86        }))
87    }
88
89    async fn handle_request_inner(&self, req: Request<Body>) -> Result<Response<Body>> {
90        // Health check endpoint
91        if req.uri().path() == "/health" && req.method() == Method::GET {
92            let request_count = self.request_count.lock().await;
93            let count = *request_count;
94            drop(request_count);
95
96            let health_data = serde_json::json!({
97                "status": "healthy",
98                "uptime": std::time::SystemTime::now()
99                    .duration_since(std::time::UNIX_EPOCH)
100                    .unwrap()
101                    .as_secs(),
102                "timestamp": chrono::Utc::now(),
103                "requestCount": count
104            });
105
106            return Ok(Response::builder()
107                .status(StatusCode::OK)
108                .header("content-type", "application/json")
109                .body(Body::from(health_data.to_string()))?);
110        }
111
112        let mut request_count = self.request_count.lock().await;
113        *request_count += 1;
114        let request_id = *request_count;
115        drop(request_count);
116
117        // Capture request body
118        let (parts, body) = req.into_parts();
119        let body_bytes = hyper::body::to_bytes(body).await?;
120        let request_body = String::from_utf8_lossy(&body_bytes);
121
122
123        self.process_request_with_config(request_body.to_string(), parts, request_id).await
124    }
125
126    async fn process_request_with_config(
127        &self,
128        request_body: String,
129        parts: http::request::Parts,
130        request_id: u64,
131    ) -> Result<Response<Body>> {
132        // Extract model name from request body
133        let mut model_name = None;
134        if !request_body.is_empty() {
135            if let Ok(parsed_body) = serde_json::from_str::<Value>(&request_body) {
136                if let Some(model) = parsed_body.get("model") {
137                    if let Some(model_str) = model.as_str() {
138                        model_name = Some(model_str.to_string());
139                    }
140                }
141            }
142        }
143
144        // Parse the target URL - handle relative URLs by prepending API base
145        let target_url = if parts.uri.path().starts_with('/') {
146            // Use model from config - always require model to be specified
147            let model = model_name.ok_or_else(|| {
148                Error::Server("Bad Request: Model must be specified in request body".to_string())
149            })?;
150            
151            let config_manager = self.config_manager.read().await;
152            let model_config = config_manager.get_model_config(&model);
153            let base_url = if model_config.api_base.ends_with('/') {
154                model_config.api_base
155            } else {
156                format!("{}/", model_config.api_base)
157            };
158            
159            
160            let path = if parts.uri.path().starts_with('/') {
161                &parts.uri.path()[1..]
162            } else {
163                parts.uri.path()
164            };
165            
166            let final_url = format!("{}{}{}", base_url, path, 
167                parts.uri.query().map_or(String::new(), |q| format!("?{}", q)));
168            
169            final_url
170        } else {
171            // Absolute URL
172            parts.uri.to_string()
173        };
174
175        // Build the proxied request
176        let reqwest_method = match parts.method.as_str() {
177            "GET" => reqwest::Method::GET,
178            "POST" => reqwest::Method::POST,
179            "PUT" => reqwest::Method::PUT,
180            "DELETE" => reqwest::Method::DELETE,
181            "PATCH" => reqwest::Method::PATCH,
182            "HEAD" => reqwest::Method::HEAD,
183            "OPTIONS" => reqwest::Method::OPTIONS,
184            _ => reqwest::Method::GET, // fallback
185        };
186
187        let mut req_builder = reqwest::Client::new()
188            .request(reqwest_method, &target_url);
189
190        // Copy headers (clean up proxy-specific ones)
191        for (name, value) in parts.headers.iter() {
192            if !matches!(name.as_str(), "host" | "proxy-connection" | "proxy-authorization") {
193                if let Ok(header_value) = value.to_str() {
194                    req_builder = req_builder.header(name.as_str(), header_value);
195                }
196            }
197        }
198
199        // Add request body
200        if !request_body.is_empty() {
201            req_builder = req_builder.body(request_body.clone());
202        }
203
204        // Make the proxied request
205        let proxy_response = req_builder.send().await?;
206        
207
208        // Check if this is an SSE response
209        let is_sse = proxy_response
210            .headers()
211            .get("content-type")
212            .and_then(|v| v.to_str().ok())
213            .map_or(false, |ct| ct.contains("text/event-stream"));
214
215        if is_sse {
216            self.handle_sse_response(proxy_response, &target_url, request_id).await
217        } else {
218            // Handle regular responses
219            let status = proxy_response.status();
220            let response_body = proxy_response.text().await?;
221            let redacted_response = self.redaction_engine.redact_sensitive_content(&response_body);
222            
223
224            Ok(Response::builder()
225                .status(status)
226                .header("content-type", "application/json")
227                .body(Body::from(redacted_response))?)
228        }
229    }
230
231    async fn handle_sse_response(
232        &self,
233        proxy_response: reqwest::Response,
234        target_url: &str,
235        request_id: u64,
236    ) -> Result<Response<Body>> {
237        use futures::StreamExt;
238        
239
240        // Initialize SSE accumulator for this request
241        self.sse_content_accumulators.write().await.insert(
242            request_id,
243            SseAccumulator {
244                buffer: String::new(),
245            },
246        );
247
248        // Detect if this is OpenAI format
249        let is_openai_format = target_url.contains("openai.com") || target_url.contains("/responses");
250
251        // Create streaming body
252        let (tx, rx) = tokio::sync::mpsc::channel::<std::result::Result<bytes::Bytes, Box<dyn std::error::Error + Send + Sync>>>(100);
253
254        // Clone necessary data for the async task
255        let redaction_engine = self.redaction_engine.clone();
256        let sse_accumulators = Arc::clone(&self.sse_content_accumulators);
257
258        // Spawn task to process SSE stream
259        tokio::spawn(async move {
260            let mut event_buffer = String::new();
261            let mut stream = proxy_response.bytes_stream();
262
263            while let Some(chunk_result) = stream.next().await {
264                match chunk_result {
265                    Ok(chunk) => {
266                        let chunk_str = String::from_utf8_lossy(&chunk);
267                        event_buffer.push_str(&chunk_str);
268
269                        // Split by double newlines to get complete events
270                        let buffer_clone = event_buffer.clone();
271                        let events: Vec<&str> = buffer_clone.split("\n\n").collect();
272                        
273                        if events.len() > 1 {
274                            // Keep the last (potentially incomplete) event in buffer
275                            event_buffer = events.last().unwrap_or(&"").to_string();
276
277                            // Process complete events
278                            for event_data in &events[..events.len() - 1] {
279                                if !event_data.trim().is_empty() {
280                                    if let Some(processed_event) = Self::process_sse_event(
281                                        event_data,
282                                        request_id,
283                                        is_openai_format,
284                                        &redaction_engine,
285                                        &sse_accumulators,
286                                    ).await {
287                                        let event_bytes = bytes::Bytes::from(format!("{}\n\n", processed_event));
288                                        if tx.send(Ok(event_bytes)).await.is_err() {
289                                            break;
290                                        }
291                                    }
292                                }
293                            }
294                        }
295                    }
296                    Err(e) => {
297                        let _ = tx.send(Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>)).await;
298                        break;
299                    }
300                }
301            }
302
303            // Handle any remaining buffered event
304            if !event_buffer.trim().is_empty() {
305                if let Some(processed_event) = Self::process_sse_event(
306                    &event_buffer,
307                    request_id,
308                    is_openai_format,
309                    &redaction_engine,
310                    &sse_accumulators,
311                ).await {
312                    let event_bytes = bytes::Bytes::from(format!("{}\n\n", processed_event));
313                    let _ = tx.send(Ok(event_bytes)).await;
314                }
315            }
316
317            // Clean up
318            sse_accumulators.write().await.remove(&request_id);
319        });
320
321        // Create streaming response
322        let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
323        let body = Body::wrap_stream(stream.map(|item| {
324            item.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
325        }));
326
327        Ok(Response::builder()
328            .status(StatusCode::OK)
329            .header("content-type", "text/event-stream")
330            .header("cache-control", "no-cache")
331            .header("connection", "keep-alive")
332            .body(body)?)
333    }
334
335    async fn process_sse_event(
336        event_data: &str,
337        request_id: u64,
338        is_openai_format: bool,
339        redaction_engine: &RedactionEngine,
340        sse_accumulators: &Arc<RwLock<HashMap<u64, SseAccumulator>>>,
341    ) -> Option<String> {
342        let lines: Vec<&str> = event_data.split('\n').collect();
343        let mut event_type = String::new();
344        let mut event_data_json: Option<serde_json::Value> = None;
345
346        // Parse SSE event
347        for line in lines {
348            if let Some(event_value) = line.strip_prefix("event: ") {
349                event_type = event_value.to_string();
350            } else if let Some(data_value) = line.strip_prefix("data: ") {
351                if let Ok(data) = serde_json::from_str::<serde_json::Value>(data_value) {
352                    event_data_json = Some(data);
353                } else {
354                    // Handle non-JSON data (like [DONE])
355                    if data_value == "[DONE]" {
356                        return Some(event_data.to_string());
357                    }
358                }
359            }
360        }
361
362        // Extract text content for redaction
363        let mut should_redact_and_forward = false;
364        let mut text_content = String::new();
365
366        if let Some(data) = &event_data_json {
367            
368            if is_openai_format {
369                // Handle OpenAI streaming format
370                if event_type == "response.output_text.delta" {
371                    if let Some(delta) = data.get("delta").and_then(|d| d.as_str()) {
372                        text_content = delta.to_string();
373                        should_redact_and_forward = true;
374                    }
375                } else if event_type == "response.output_text.done" {
376                    if let Some(text) = data.get("text").and_then(|t| t.as_str()) {
377                        text_content = text.to_string();
378                        should_redact_and_forward = true;
379                    }
380                }
381            } else {
382                // Handle Claude streaming format  
383                if event_type == "content_block_delta" {
384                    if let Some(delta) = data.get("delta") {
385                        if let Some(text) = delta.get("text").and_then(|t| t.as_str()) {
386                            text_content = text.to_string();
387                            should_redact_and_forward = true;
388                        }
389                    }
390                }
391            }
392            
393            if !should_redact_and_forward {
394            }
395        }
396
397        if should_redact_and_forward && !text_content.is_empty() {
398            
399            // Apply redaction with buffering (same as Node.js)
400            let redacted_content = if is_openai_format && event_type == "response.output_text.done" {
401                // For complete text events, apply redaction directly
402                let redacted = redaction_engine.redact_sensitive_content(&text_content);
403                Some(redacted)
404            } else {
405                // For streaming deltas, use buffering
406                let result = Self::process_chunk_with_buffer(&text_content, request_id, redaction_engine, sse_accumulators).await;
407                result
408            };
409
410            if let Some(redacted) = redacted_content {
411                if let Some(mut data) = event_data_json {
412                    if is_openai_format {
413                        // Update OpenAI format
414                        if event_type == "response.output_text.delta" {
415                            data["delta"] = serde_json::Value::String(redacted);
416                        } else if event_type == "response.output_text.done" {
417                            data["text"] = serde_json::Value::String(redacted);
418                        }
419                    } else {
420                        // Update Claude format
421                        if let Some(delta) = data.get_mut("delta") {
422                            delta["text"] = serde_json::Value::String(redacted);
423                        }
424                    }
425
426                    return Some(format!("event: {}\ndata: {}", event_type, data));
427                }
428            } else {
429                // Don't forward this event since we're buffering
430                return None;
431            }
432        } else if event_type == "content_block_stop" && !is_openai_format {
433            // Claude-specific: Flush any remaining buffer content
434            if let Some(final_chunk) = Self::flush_buffer(request_id, redaction_engine, sse_accumulators).await {
435                let final_event_data = serde_json::json!({
436                    "type": "content_block_delta",
437                    "index": 0,
438                    "delta": {
439                        "type": "text_delta",
440                        "text": final_chunk
441                    }
442                });
443                return Some(format!("event: content_block_delta\ndata: {}\n\nevent: content_block_stop\ndata: {}", 
444                    final_event_data, event_data_json.unwrap_or(serde_json::Value::Null)));
445            }
446        } else if is_openai_format && (event_type == "response.completed" || event_data.contains("[DONE]")) {
447            // OpenAI-specific: Handle end of stream and flush buffer
448            if let Some(final_chunk) = Self::flush_buffer(request_id, redaction_engine, sse_accumulators).await {
449                let final_data = serde_json::json!({
450                    "type": "response.output_text.delta",
451                    "delta": final_chunk
452                });
453                return Some(format!("event: response.output_text.delta\ndata: {}\n\n{}", 
454                    final_data, event_data));
455            }
456        }
457
458        // Forward all other events as-is
459        Some(event_data.to_string())
460    }
461
462    async fn process_chunk_with_buffer(
463        new_chunk: &str,
464        request_id: u64,
465        redaction_engine: &RedactionEngine,
466        sse_accumulators: &Arc<RwLock<HashMap<u64, SseAccumulator>>>,
467    ) -> Option<String> {
468        let mut accumulators = sse_accumulators.write().await;
469        if let Some(accumulator) = accumulators.get_mut(&request_id) {
470            // Add new chunk to buffer
471            accumulator.buffer.push_str(new_chunk);
472
473            // Split by lines and process complete lines
474            let buffer_copy = accumulator.buffer.clone();
475            let mut lines: Vec<&str> = buffer_copy.split('\n').collect();
476            
477            // Keep the last (potentially incomplete) line in buffer (like Node.js pop())
478            accumulator.buffer = lines.pop().unwrap_or("").to_string();
479            
480            // Process complete lines
481            if !lines.is_empty() {
482                let complete_lines = lines.join("\n") + "\n";
483                let redacted = redaction_engine.redact_sensitive_content(&complete_lines);
484                return Some(redacted);
485            }
486            
487            // No complete lines yet, don't send anything
488            return None;
489        }
490        None
491    }
492
493    async fn flush_buffer(
494        request_id: u64,
495        redaction_engine: &RedactionEngine,
496        sse_accumulators: &Arc<RwLock<HashMap<u64, SseAccumulator>>>,
497    ) -> Option<String> {
498        let mut accumulators = sse_accumulators.write().await;
499        if let Some(accumulator) = accumulators.get_mut(&request_id) {
500            if accumulator.buffer.is_empty() {
501                return None;
502            }
503            
504            // Process remaining buffer content
505            let remaining = accumulator.buffer.clone();
506            accumulator.buffer.clear();
507            
508            let redacted = redaction_engine.redact_sensitive_content(&remaining);
509            Some(redacted)
510        } else {
511            None
512        }
513    }
514}
515
516impl Clone for ProxyServer {
517    fn clone(&self) -> Self {
518        let https = HttpsConnector::new();
519        let client = Client::builder().build::<_, hyper::Body>(https);
520        
521        Self {
522            port: self.port,
523            request_count: Arc::clone(&self.request_count),
524            redaction_engine: self.redaction_engine.clone(),
525            config_manager: Arc::clone(&self.config_manager),
526            client,
527            sse_content_accumulators: Arc::clone(&self.sse_content_accumulators),
528        }
529    }
530}
531
532impl Clone for RedactionEngine {
533    fn clone(&self) -> Self {
534        RedactionEngine::new()
535    }
536}
vibekit_proxy/server.rs

vibekit_proxy/
server.rs
