asimov-telemetry 25.9.0

// This is free and unencumbered software released into the public domain.

use crate::{Event, Outcome};
use alloc::{
    collections::BTreeMap,
    format,
    string::{String, ToString},
    vec::Vec,
};
use asimov_http::proxy::{OriginProxy, ProxyOptions};
use core::time::Duration;
use reqwest::header::{HeaderMap, HeaderValue};
use serde::Serialize;
use serde_json::{Value, json};
use std::{
    fs::{self, File},
    io::{self, BufRead, BufReader, Write},
    path::Path,
    sync::{Mutex, mpsc},
    time::{SystemTime, UNIX_EPOCH},
};

const ENDPOINT: &str = "https://api.statsig.com/v1/log_event";
const FLUSH_INTERVAL: Duration = Duration::from_secs(5 * 60);
const STALE_AFTER: Duration = Duration::from_secs(60 * 60);
const STALE_SHUTDOWN_WAIT: Duration = Duration::from_millis(500);
const MAX_EVENT_AGE: Duration = Duration::from_secs(30 * 24 * 60 * 60);
const BATCH_SIZE: usize = 500;

pub struct Telemetry {
    events: Mutex<File>,
    user: User,
    invocation_id: String,
    caller: Caller,
    stale_flush: Mutex<Option<mpsc::Receiver<()>>>,
}

impl Telemetry {
    /// Events are appended to a per-invocation log in `directory`, which stays
    /// locked while the process lives. Logs of exited processes are sent at most
    /// once every five minutes.
    ///
    /// Uploads use shared environment/native proxy discovery from `asimov-http`.
    /// Redirects are not followed; failed uploads remain queued for a later flush.
    pub fn new(client_key: &str, app_version: &str, directory: &Path) -> Option<Self> {
        if is_disabled(directory) {
            return None;
        }
        Self::start(client_key, app_version, directory, ENDPOINT)
    }

    fn start(
        client_key: &str,
        app_version: &str,
        directory: &Path,
        endpoint: &str,
    ) -> Option<Self> {
        if !client_key.starts_with("client-") || client_key.len() <= "client-".len() {
            return None;
        }
        let mut key = HeaderValue::from_str(client_key).ok()?;
        key.set_sensitive(true);
        fs::create_dir_all(directory).ok()?;
        let id_path = directory.join("id");
        let new_install = !id_path.exists();
        if new_install {
            let id = bs58::encode(dogma_uuid::Uuid::new_v4().as_bytes()).into_string();
            fs::write(&id_path, id).ok()?;
        }
        let user_id = fs::read_to_string(id_path).ok()?;

        let invocation_id = dogma_uuid::Uuid::now_v7().to_string();
        let events =
            File::create_new(directory.join(format!("events-{invocation_id}.jsonl"))).ok()?;
        events.try_lock().ok()?;

        let user = User {
            user_id,
            custom_ids: fs::read_to_string(directory.join("custom-ids.json"))
                .ok()
                .and_then(|ids| serde_json::from_str(&ids).ok())
                .unwrap_or_default(),
            app_version: app_version.to_string(),
            custom: BTreeMap::from([
                ("target_os", std::env::consts::OS),
                ("target_arch", std::env::consts::ARCH),
            ]),
        };
        let caller = Caller::detect();
        if new_install {
            let event = wire_event(
                Event::Installed,
                timestamp(),
                &user,
                &invocation_id,
                &caller,
            );
            // On its own thread, as the caller may be inside an async runtime.
            let sent = std::thread::scope(|scope| {
                scope
                    .spawn(|| send(key.clone(), endpoint, core::slice::from_ref(&event)))
                    .join()
            });
            if !matches!(sent, Ok(Some(()))) {
                let _ = writeln!(&events, "{event}");
            }
        }

        let age = |name: &str| {
            fs::metadata(directory.join(name))
                .and_then(|metadata| metadata.modified())
                .map(|time| SystemTime::now().duration_since(time).unwrap_or_default())
                .unwrap_or(Duration::MAX)
        };
        let flush = (age("attempted") >= FLUSH_INTERVAL)
            .then(|| {
                let (sender, receiver) = mpsc::channel();
                let directory = directory.to_path_buf();
                let endpoint = endpoint.to_string();
                std::thread::Builder::new()
                    .name("asimov-telemetry".into())
                    .spawn(move || {
                        flush(&directory, key, &endpoint);
                        let _ = sender.send(());
                    })
                    .ok()
                    .map(|_| receiver)
            })
            .flatten();

        Some(Self {
            events: Mutex::new(events),
            user,
            invocation_id,
            caller,
            stale_flush: Mutex::new(flush.filter(|_| age("flushed") >= STALE_AFTER)),
        })
    }

    pub fn log(&self, event: Event) {
        let event = wire_event(
            event,
            timestamp(),
            &self.user,
            &self.invocation_id,
            &self.caller,
        );
        if let Ok(mut events) = self.events.lock() {
            let _ = writeln!(events, "{event}");
        }
    }

    /// Only waits for an in-flight flush when the last successful one is stale.
    pub fn shutdown(&self) {
        let flush = self
            .stale_flush
            .lock()
            .ok()
            .and_then(|mut flush| flush.take());
        if let Some(flush) = flush {
            let _ = flush.recv_timeout(STALE_SHUTDOWN_WAIT);
        }
    }
}

/// Whether `ASIMOV_TELEMETRY` is `0`, `false`, `no` or `off`, `DO_NOT_TRACK` is `1`, or
/// [`disable`] was called for `directory`.
pub fn is_disabled(directory: &Path) -> bool {
    std::env::var("ASIMOV_TELEMETRY").is_ok_and(|value| {
        matches!(
            value.trim().to_ascii_lowercase().as_str(),
            "0" | "false" | "no" | "off"
        )
    }) || std::env::var("DO_NOT_TRACK").is_ok_and(|value| value.trim() == "1")
        || directory.join("disabled").exists()
}

pub fn enable(directory: &Path) -> io::Result<()> {
    match fs::remove_file(directory.join("disabled")) {
        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
        result => result,
    }
}

/// Also discards any unsent events.
pub fn disable(directory: &Path) -> io::Result<()> {
    fs::create_dir_all(directory)?;
    File::create(directory.join("disabled"))?;
    for entry in fs::read_dir(directory)?.flatten() {
        if entry.file_name().to_string_lossy().starts_with("events-") {
            fs::remove_file(entry.path())?;
        }
    }
    Ok(())
}

fn flush(directory: &Path, key: HeaderValue, endpoint: &str) -> Option<()> {
    let touch = |name: &str| {
        File::options()
            .create(true)
            .truncate(false)
            .write(true)
            .open(directory.join(name))
            .and_then(|file| file.set_modified(SystemTime::now()))
            .ok()
    };
    let lock = File::create(directory.join("flush.lock")).ok()?;
    lock.try_lock().ok()?;
    touch("attempted")?;

    let mut logs = Vec::new();
    for path in fs::read_dir(directory)
        .ok()?
        .flatten()
        .map(|entry| entry.path())
    {
        let is_log = path
            .file_name()
            .and_then(|name| name.to_str())
            .is_some_and(|name| name.starts_with("events-") && name.ends_with(".jsonl"));
        let Some(file) = is_log.then(|| File::open(&path).ok()).flatten() else {
            continue;
        };
        if file.try_lock().is_err() {
            continue;
        }
        let expired = file
            .metadata()
            .and_then(|metadata| metadata.modified())
            .is_ok_and(|time| time.elapsed().is_ok_and(|age| age > MAX_EVENT_AGE));
        if expired {
            drop(file);
            let _ = fs::remove_file(&path);
            continue;
        }
        logs.push((path, file));
    }

    let events: Vec<Value> = logs
        .iter()
        .flat_map(|(_, file)| BufReader::new(file).lines().map_while(Result::ok))
        .filter_map(|line| serde_json::from_str(&line).ok())
        .collect();
    if !events.is_empty() {
        send(key, endpoint, &events)?;
    }

    for (path, file) in logs {
        drop(file);
        let _ = fs::remove_file(path);
    }
    touch("flushed")
}

fn send(key: HeaderValue, endpoint: &str, events: &[Value]) -> Option<()> {
    let endpoint = reqwest::Url::parse(endpoint).ok()?;
    let proxy = OriginProxy::resolve(&endpoint, &ProxyOptions::default()).ok()?;
    let mut headers = HeaderMap::new();
    headers.insert("statsig-api-key", key);
    headers.insert("statsig-sdk-type", HeaderValue::from_static("asimov.rs"));
    headers.insert(
        "statsig-sdk-version",
        HeaderValue::from_static(env!("CARGO_PKG_VERSION")),
    );
    let builder = reqwest::Client::builder()
        .default_headers(headers)
        .timeout(Duration::from_secs(10));
    let client = proxy.configure_reqwest(builder).ok()?.build().ok()?;
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .ok()?;
    runtime.block_on(async {
        for batch in events.chunks(BATCH_SIZE) {
            // Every batch uses the exact URL whose proxy policy was resolved.
            let response = client
                .post(endpoint.clone())
                .header("statsig-client-time", timestamp().to_string())
                .json(&json!({"events": batch}))
                .send()
                .await
                .ok()?;
            if !response.status().is_success() {
                return None;
            }
        }
        Some(())
    })
}

fn timestamp() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis()
}

struct Caller {
    kind: &'static str,
    agent: &'static str,
    source: &'static str,
}

impl Caller {
    fn detect() -> Self {
        let Some(agent) = agent_detector::detect() else {
            return Self {
                kind: "unknown",
                agent: "unknown",
                source: "none",
            };
        };
        Self {
            kind: "agent",
            agent: Self::agent_name(&agent.name),
            source: match agent.source {
                agent_detector::DetectionSource::ParentProcess => "process",
                agent_detector::DetectionSource::StandardEnvVar
                | agent_detector::DetectionSource::ToolEnvVar => "environment",
                _ => "unknown",
            },
        }
    }

    fn agent_name(name: &str) -> &'static str {
        // `AI_AGENT` may carry a suffix, e.g. Claude Code sets `claude-code_2-1-286_agent`.
        [
            "pi",
            "hermes",
            "claude-code",
            "codex",
            "openclaw",
            "opencode",
            "cursor",
            "cursor-cli",
            "gemini",
            "cowork",
            "aider",
            "github-copilot",
            "goose",
            "cline",
            "oh-my-pi",
        ]
        .into_iter()
        .find(|known| {
            name.strip_prefix(known).is_some_and(|suffix| {
                !suffix.starts_with(|c: char| c.is_ascii_alphanumeric() || c == '-')
            })
        })
        .unwrap_or("other")
    }
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct User {
    #[serde(rename = "userID")]
    user_id: String,
    #[serde(rename = "customIDs", skip_serializing_if = "BTreeMap::is_empty")]
    custom_ids: BTreeMap<String, String>,
    app_version: String,
    custom: BTreeMap<&'static str, &'static str>,
}

fn wire_event(
    event: Event,
    time: u128,
    user: &User,
    invocation_id: &str,
    caller: &Caller,
) -> Value {
    let mut metadata = BTreeMap::from([
        ("schema_version", "1".to_string()),
        ("invocation_id", invocation_id.to_string()),
        ("caller_kind", caller.kind.to_string()),
        ("agent_name", caller.agent.to_string()),
        ("detection_source", caller.source.to_string()),
    ]);
    let name = match event {
        Event::Installed => "cli_installed",
        Event::CommandStarted { command, modules } => {
            metadata.insert("command", command);
            if let [module] = modules.as_slice() {
                metadata.insert("module", module.clone());
            }
            metadata.insert("modules", json!(modules).to_string());
            "cli_command_started"
        },
        Event::CommandFinished {
            command,
            modules,
            exit_code,
            duration,
        } => {
            metadata.insert("command", command);
            if let [module] = modules.as_slice() {
                metadata.insert("module", module.clone());
            }
            metadata.insert("modules", json!(modules).to_string());
            metadata.insert("exit_code", exit_code.to_string());
            metadata.insert(
                "outcome",
                if exit_code == 0 { "success" } else { "failure" }.to_string(),
            );
            metadata.insert("duration_ms", duration.as_millis().to_string());
            "cli_command_finished"
        },
        Event::ModuleOperationFinished {
            operation,
            module,
            outcome,
            duration,
        } => {
            metadata.insert(
                "operation",
                match operation {
                    crate::Operation::Fetch => "fetch",
                    crate::Operation::Read => "read",
                    crate::Operation::List => "list",
                }
                .to_string(),
            );
            metadata.insert("module", module);
            metadata.insert(
                "outcome",
                match outcome {
                    Outcome::Success => "success",
                    Outcome::Failure => "failure",
                }
                .to_string(),
            );
            metadata.insert("duration_ms", duration.as_millis().to_string());
            "cli_module_operation_finished"
        },
    };
    json!({"eventName": name, "user": user, "time": time, "metadata": metadata})
}

#[cfg(test)]
mod tests;