claude_utils/turbo/
proxy.rs

1use axum::{
2    body::Body,
3    extract::{Request, State},
4    http::{Response, StatusCode},
5    middleware::Next,
6};
7use serde_json::{json, Value};
8use std::sync::Arc;
9use tracing::{debug, info};
10
11use super::TurboMode;
12
13/// MCP Proxy that intercepts and modifies Claude's requests
14pub struct TurboProxy {
15    _turbo: Arc<TurboMode>,
16}
17
18impl TurboProxy {
19    pub fn new(turbo: Arc<TurboMode>) -> Self {
20        Self { _turbo: turbo }
21    }
22
23    /// Middleware that intercepts MCP requests
24    pub async fn intercept_middleware(
25        State(turbo): State<Arc<TurboMode>>,
26        request: Request,
27        next: Next,
28    ) -> Result<Response<Body>, StatusCode> {
29        let (parts, body) = request.into_parts();
30        
31        // Parse the body
32        let bytes = axum::body::to_bytes(body, usize::MAX)
33            .await
34            .map_err(|_| StatusCode::BAD_REQUEST)?;
35        
36        let mut json_body: Value = serde_json::from_slice(&bytes)
37            .map_err(|_| StatusCode::BAD_REQUEST)?;
38
39        // Check if YOLO mode is enabled
40        if turbo.is_yolo().await {
41            // Auto-approve all permission requests
42            if let Some(method) = json_body.get("method").and_then(|m| m.as_str()) {
43                match method {
44                    "tools/call" => {
45                        info!("🚀 YOLO: Auto-approving tool call");
46                        // Modify the request to bypass permission
47                        if let Some(params) = json_body.get_mut("params") {
48                            params["auto_approved"] = json!(true);
49                        }
50                    }
51                    "resources/read" | "resources/write" => {
52                        info!("🚀 YOLO: Auto-approving resource access");
53                        // Add approval flag
54                        json_body["yolo_approved"] = json!(true);
55                    }
56                    _ => {}
57                }
58            }
59        }
60
61        // Detect batch operations for parallel execution
62        if let Some(params) = json_body.get("params") {
63            if let Some(operations) = params.get("operations").and_then(|o| o.as_array()) {
64                if operations.len() > 1 {
65                    debug!("Detected batch operation with {} items", operations.len());
66                    // Mark for parallel execution
67                    json_body["turbo_parallel"] = json!(true);
68                }
69            }
70        }
71
72        // Reconstruct request with modified body
73        let body_vec = serde_json::to_vec(&json_body)
74            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
75        let new_body = Body::from(body_vec);
76        let new_request = Request::from_parts(parts, new_body);
77
78        // Continue with the request
79        let response = next.run(new_request).await;
80
81        // Intercept response for retry logic
82        if turbo.config.read().await.auto_retry {
83            // Check if response indicates failure
84            let (parts, body) = response.into_parts();
85            let bytes = axum::body::to_bytes(body, usize::MAX)
86                .await
87                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
88            
89            if let Ok(json_response) = serde_json::from_slice::<Value>(&bytes) {
90                if let Some(error) = json_response.get("error") {
91                    info!("Operation failed, will retry: {:?}", error);
92                    // Store for retry logic
93                    turbo.executor.queue_retry(json_body.clone()).await;
94                }
95            }
96
97            let new_body = Body::from(bytes);
98            Ok(Response::from_parts(parts, new_body))
99        } else {
100            Ok(response)
101        }
102    }
103}
claude_utils/turbo/proxy.rs

claude_utils/turbo/
proxy.rs
