claude-utils 0.2.0

Cross-platform companion toolkit for Anthropic's Claude Code CLI
Documentation
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;

/// MCP Proxy that intercepts and modifies Claude's requests
pub struct TurboProxy {
    _turbo: Arc<TurboMode>,
}

impl TurboProxy {
    pub fn new(turbo: Arc<TurboMode>) -> Self {
        Self { _turbo: turbo }
    }

    /// Middleware that intercepts MCP requests
    pub async fn intercept_middleware(
        State(turbo): State<Arc<TurboMode>>,
        request: Request,
        next: Next,
    ) -> Result<Response<Body>, StatusCode> {
        let (parts, body) = request.into_parts();
        
        // Parse the body
        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)?;

        // Check if YOLO mode is enabled
        if turbo.is_yolo().await {
            // Auto-approve all permission requests
            if let Some(method) = json_body.get("method").and_then(|m| m.as_str()) {
                match method {
                    "tools/call" => {
                        info!("🚀 YOLO: Auto-approving tool call");
                        // Modify the request to bypass permission
                        if let Some(params) = json_body.get_mut("params") {
                            params["auto_approved"] = json!(true);
                        }
                    }
                    "resources/read" | "resources/write" => {
                        info!("🚀 YOLO: Auto-approving resource access");
                        // Add approval flag
                        json_body["yolo_approved"] = json!(true);
                    }
                    _ => {}
                }
            }
        }

        // Detect batch operations for parallel execution
        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());
                    // Mark for parallel execution
                    json_body["turbo_parallel"] = json!(true);
                }
            }
        }

        // Reconstruct request with modified body
        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);

        // Continue with the request
        let response = next.run(new_request).await;

        // Intercept response for retry logic
        if turbo.config.read().await.auto_retry {
            // Check if response indicates failure
            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);
                    // Store for retry logic
                    turbo.executor.queue_retry(json_body.clone()).await;
                }
            }

            let new_body = Body::from(bytes);
            Ok(Response::from_parts(parts, new_body))
        } else {
            Ok(response)
        }
    }
}