mcp-memory/server/src/api/telemetry.rs

use crate::AppState;
use axum::extract::State;
use std::sync::Arc;

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct NvimTelemetry {
    pub session_id: String,
    pub event: String,
    pub file: Option<String>,
    pub line: Option<i64>,
    pub col: Option<i64>,
    // New interactive fields
    pub prompt: Option<String>,
    pub response: Option<String>,
    pub index: Option<i64>,
    pub request_id: Option<String>,
    pub diff_id: Option<String>,
    pub start_row: Option<i64>,
    pub end_row: Option<i64>,
}

pub async fn nvim_telemetry_handler(
    State(state): State<Arc<AppState>>,
    axum::Json(payload): axum::Json<NvimTelemetry>,
) -> impl axum::response::IntoResponse {
    // 1. Update active_nvim.txt if FocusGained, BufEnter, or VimEnter
    if payload.event == "FocusGained" || payload.event == "BufEnter" || payload.event == "VimEnter"
    {
        let session = &payload.session_id;
        let is_unix_socket = session.starts_with('/') || session.starts_with('~');

        if is_unix_socket {
            let wsl_path = "\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
            let _ = tokio::fs::write(wsl_path, session).await;
        } else {
            let profile =
                std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
            let win_path = format!("{}\\.gemini\\active_nvim.txt", profile);
            let _ = tokio::fs::write(&win_path, session).await;
        }
    }

    // 2. Broadcast to UI WebSockets
    let ws_msg = serde_json::json!({
        "type": "nvim_telemetry",
        "data": payload
    });

    let msg_str = ws_msg.to_string();
    let senders: Vec<_> = state
        .clients
        .read()
        .unwrap_or_else(|e| e.into_inner())
        .values()
        .cloned()
        .collect();
    for tx in senders {
        let _ = tx.try_send(msg_str.clone());
    }

    if payload.event == "BufWritePost"
        && let Some(ref file_path) = payload.file
    {
        let normalized_file = file_path.replace("\\", "/");
        let topic = format!("nvim:save:{}", normalized_file);
        let event = crate::state::GenericEvent {
            topic,
            session_id: Some(payload.session_id.clone()),
            payload: serde_json::json!(&payload),
        };
        let _ = state.handler.state.event_bus_tx.send(event);
    }

    if let Some(ref f) = payload.file {
        state.handler.state.record_activity(
            "editor",
            &format!("Neovim {}: {}", payload.event, f),
            None,
        );
    }

    axum::Json(serde_json::json!({"status": "ok"}))
}

use crate::models::TerminalHistory;

pub async fn get_terminal_history_handler(
    State(state): State<Arc<AppState>>,
) -> impl axum::response::IntoResponse {
    let history_json = state
        .handler
        .state
        .telemetry
        .terminal_history
        .read_with(|h| serde_json::to_string(h).unwrap_or_else(|_| "[]".to_string()));
    (
        [(axum::http::header::CONTENT_TYPE, "application/json")],
        history_json,
    )
}

pub async fn terminal_telemetry_handler(
    State(state): State<Arc<AppState>>,
    axum::Json(payload): axum::Json<TerminalHistory>,
) -> impl axum::response::IntoResponse {
    state.handler.state.record_activity(
        "terminal",
        &format!("Ran command: {}", payload.command),
        Some(&format!(
            "exit_code: {}, cwd: {}",
            payload.exit_code, payload.cwd
        )),
    );

    state.handler.state.record_terminal_history(payload.clone());

    let ws_msg = serde_json::json!({
        "type": "terminal_telemetry",
        "data": payload
    });

    let msg_str = ws_msg.to_string();
    let senders: Vec<_> = state
        .clients
        .read()
        .unwrap_or_else(|e| e.into_inner())
        .values()
        .cloned()
        .collect();
    for tx in senders {
        let _ = tx.try_send(msg_str.clone());
    }

    axum::Json(serde_json::json!({"status": "ok"}))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::router::MemoryHandler;
    use crate::state::MemoryState;
    use axum::http::Request;
    use axum::routing::get;
    use axum::routing::post;
    use std::collections::HashMap;
    use std::sync::atomic::AtomicUsize;
    use tempfile::tempdir;
    use tower::ServiceExt;

    #[tokio::test]
    async fn test_terminal_history() {
        let dir = tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
        let (shutdown_tx, _) = tokio::sync::oneshot::channel();
        let app_state = Arc::new(AppState {
            handler: Arc::new(MemoryHandler::new(state.clone())),
            clients: std::sync::RwLock::new(HashMap::new()),
            next_id: AtomicUsize::new(1),
            shutdown_tx: std::sync::Mutex::new(Some(shutdown_tx)),
        });

        let app = axum::Router::new()
            .route("/api/terminal/history", get(get_terminal_history_handler))
            .route("/terminal/telemetry", post(terminal_telemetry_handler))
            .with_state(app_state.clone());

        // Create a fake terminal telemetry payload
        let payload = TerminalHistory {
            command: "echo test".to_string(),
            exit_code: 0,
            cwd: "C:/Users/reazul.ashraf/workspace".to_string(),
            os: "windows".to_string(),
            timestamp: 12345,
            error_output: None,
            ..Default::default()
        };

        // Post the telemetry
        let req = Request::builder()
            .method("POST")
            .uri("/terminal/telemetry")
            .header("content-type", "application/json")
            .body(axum::body::Body::from(
                serde_json::to_string(&payload).unwrap(),
            ))
            .unwrap();

        let response = app.clone().oneshot(req).await.unwrap();
        assert_eq!(response.status(), 200);

        // Get the telemetry
        let req2 = Request::builder()
            .method("GET")
            .uri("/api/terminal/history")
            .body(axum::body::Body::empty())
            .unwrap();

        let response2 = app.oneshot(req2).await.unwrap();
        assert_eq!(response2.status(), 200);

        let body_bytes = axum::body::to_bytes(response2.into_body(), usize::MAX)
            .await
            .unwrap();
        let history: Vec<TerminalHistory> = serde_json::from_slice(&body_bytes).unwrap();
        assert_eq!(history.len(), 1);
        assert_eq!(history[0].command, "echo test");
    }
}