use axum::{
body::Body,
extract::{Request, State},
http::{Response, StatusCode},
middleware::Next,
};
use serde_json::{json, Value};
use std::sync::Arc;
use tracing::{debug, info};
use super::TurboMode;
pub struct TurboProxy {
_turbo: Arc<TurboMode>,
}
impl TurboProxy {
pub fn new(turbo: Arc<TurboMode>) -> Self {
Self { _turbo: turbo }
}
pub async fn intercept_middleware(
State(turbo): State<Arc<TurboMode>>,
request: Request,
next: Next,
) -> Result<Response<Body>, StatusCode> {
let (parts, body) = request.into_parts();
let bytes = axum::body::to_bytes(body, usize::MAX)
.await
.map_err(|_| StatusCode::BAD_REQUEST)?;
let mut json_body: Value = serde_json::from_slice(&bytes)
.map_err(|_| StatusCode::BAD_REQUEST)?;
if turbo.is_yolo().await {
if let Some(method) = json_body.get("method").and_then(|m| m.as_str()) {
match method {
"tools/call" => {
info!("🚀 YOLO: Auto-approving tool call");
if let Some(params) = json_body.get_mut("params") {
params["auto_approved"] = json!(true);
}
}
"resources/read" | "resources/write" => {
info!("🚀 YOLO: Auto-approving resource access");
json_body["yolo_approved"] = json!(true);
}
_ => {}
}
}
}
if let Some(params) = json_body.get("params") {
if let Some(operations) = params.get("operations").and_then(|o| o.as_array()) {
if operations.len() > 1 {
debug!("Detected batch operation with {} items", operations.len());
json_body["turbo_parallel"] = json!(true);
}
}
}
let body_vec = serde_json::to_vec(&json_body)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let new_body = Body::from(body_vec);
let new_request = Request::from_parts(parts, new_body);
let response = next.run(new_request).await;
if turbo.config.read().await.auto_retry {
let (parts, body) = response.into_parts();
let bytes = axum::body::to_bytes(body, usize::MAX)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
if let Ok(json_response) = serde_json::from_slice::<Value>(&bytes) {
if let Some(error) = json_response.get("error") {
info!("Operation failed, will retry: {:?}", error);
turbo.executor.queue_retry(json_body.clone()).await;
}
}
let new_body = Body::from(bytes);
Ok(Response::from_parts(parts, new_body))
} else {
Ok(response)
}
}
}