feat: initial commit — antigravity proxy with MITM, standalone LS, and snapshot tooling

src/api/mod.rs

@@ -0,0 +1,176 @@
+//! Axum API server — OpenAI-compatible Responses + Chat Completions endpoints.
+
+mod completions;
+mod models;
+mod polling;
+mod responses;
+mod types;
+mod util;
+
+use crate::constants::safe_truncate;
+use crate::session::SessionManager;
+use axum::{
+    extract::{DefaultBodyLimit, State},
+    http::StatusCode,
+    response::{IntoResponse, Json},
+    routing::{delete, get, post},
+    Router,
+};
+use std::sync::Arc;
+use tower_http::cors::CorsLayer;
+use tracing::warn;
+
+use self::models::MODELS;
+use self::types::TokenRequest;
+
+// ─── Shared state ────────────────────────────────────────────────────────────
+
+pub struct AppState {
+    pub backend: Arc<crate::backend::Backend>,
+    pub sessions: SessionManager,
+    pub mitm_store: crate::mitm::store::MitmStore,
+    pub quota_store: crate::quota::QuotaStore,
+}
+
+// ─── Router ──────────────────────────────────────────────────────────────────
+
+pub fn router(state: Arc<AppState>) -> Router {
+    Router::new()
+        .route("/v1/responses", post(responses::handle_responses))
+        .route(
+            "/v1/chat/completions",
+            post(completions::handle_completions),
+        )
+        .route("/v1/models", get(handle_models))
+        .route("/v1/sessions", get(handle_list_sessions))
+        .route("/v1/sessions/{id}", delete(handle_delete_session))
+        .route("/v1/token", post(handle_set_token))
+        .route("/v1/usage", get(handle_usage))
+        .route("/v1/quota", get(handle_quota))
+        .route("/health", get(handle_health))
+        .route("/", get(handle_root))
+        .layer(CorsLayer::permissive())
+        .layer(DefaultBodyLimit::max(1_048_576)) // 1 MB
+        .with_state(state)
+}
+
+// ─── Simple handlers ─────────────────────────────────────────────────────────
+
+async fn handle_root() -> Json<serde_json::Value> {
+    Json(serde_json::json!({
+        "service": "antigravity-openai-proxy",
+        "version": "3.2.0",
+        "runtime": "rust",
+        "endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses",
+            "/v1/models",
+            "/v1/sessions",
+            "/v1/token",
+            "/v1/usage",
+            "/v1/quota",
+            "/health",
+        ],
+    }))
+}
+
+async fn handle_health() -> Json<serde_json::Value> {
+    Json(serde_json::json!({"status": "ok"}))
+}
+
+async fn handle_models() -> Json<serde_json::Value> {
+    let models: Vec<serde_json::Value> = MODELS
+        .iter()
+        .map(|m| {
+            serde_json::json!({
+                "id": m.name,
+                "object": "model",
+                "created": 1700000000u64,
+                "owned_by": "antigravity",
+                "permission": [],
+                "root": m.name,
+                "parent": null,
+                "meta": {
+                    "label": m.label,
+                    "enum_value": m.model_enum,
+                },
+            })
+        })
+        .collect();
+    Json(serde_json::json!({"object": "list", "data": models}))
+}
+
+async fn handle_list_sessions(
+    State(state): State<Arc<AppState>>,
+) -> Json<serde_json::Value> {
+    let sessions = state.sessions.list_sessions().await;
+    Json(serde_json::json!({"sessions": sessions}))
+}
+
+async fn handle_delete_session(
+    State(state): State<Arc<AppState>>,
+    axum::extract::Path(id): axum::extract::Path<String>,
+) -> impl IntoResponse {
+    if state.sessions.delete_session(&id).await {
+        (
+            StatusCode::OK,
+            Json(serde_json::json!({"status": "deleted", "session_id": id})),
+        )
+    } else {
+        (
+            StatusCode::NOT_FOUND,
+            Json(serde_json::json!({"error": format!("Session not found: {id}")})),
+        )
+    }
+}
+
+async fn handle_set_token(
+    State(state): State<Arc<AppState>>,
+    Json(body): Json<TokenRequest>,
+) -> impl IntoResponse {
+    if !body.token.starts_with("ya29.") {
+        return (
+            StatusCode::BAD_REQUEST,
+            Json(serde_json::json!({"error": "Invalid token. Must start with ya29."})),
+        );
+    }
+    state.backend.set_oauth_token(body.token.clone()).await;
+
+    // Also persist to file
+    let token_path = crate::constants::token_file_path();
+    if let Err(e) = tokio::fs::write(&token_path, &body.token).await {
+        warn!("Failed to write token file: {e}");
+    }
+
+    let preview = safe_truncate(&body.token, 20);
+    (
+        StatusCode::OK,
+        Json(serde_json::json!({"status": "ok", "token_prefix": preview})),
+    )
+}
+
+async fn handle_usage(
+    State(state): State<Arc<AppState>>,
+) -> Json<serde_json::Value> {
+    let stats = state.mitm_store.stats().await;
+    Json(serde_json::json!({
+        "mitm": {
+            "total_requests": stats.total_requests,
+            "total_input_tokens": stats.total_input_tokens,
+            "total_output_tokens": stats.total_output_tokens,
+            "total_cache_read_tokens": stats.total_cache_read_tokens,
+            "total_cache_creation_tokens": stats.total_cache_creation_tokens,
+            "total_thinking_output_tokens": stats.total_thinking_output_tokens,
+            "total_response_output_tokens": stats.total_response_output_tokens,
+            "total_tokens": stats.total_input_tokens + stats.total_output_tokens,
+            "per_model": stats.per_model,
+        }
+    }))
+}
+
+async fn handle_quota(
+    State(state): State<Arc<AppState>>,
+) -> Json<serde_json::Value> {
+    let snap = state.quota_store.snapshot().await;
+    Json(serde_json::to_value(snap).unwrap_or_default())
+}