oc_collect/

lib.rs

use serde::Deserialize;
use sqlx::SqlitePool;

#[derive(Deserialize, Debug)]
pub struct Model {
    pub id: String,
    #[serde(rename = "providerID")]
    pub provider_id: String,
    pub variant: Option<String>,
}

#[derive(Debug, sqlx::FromRow)]
pub struct RowUsageSession {
    pub id: String,
    pub project_id: String,
    pub slug: String,
    pub title: String,
    pub cost: f64,
    pub tokens_input: i64,
    pub tokens_output: i64,
    pub tokens_reasoning: i64,
    pub tokens_cache_read: i64,
    pub tokens_cache_write: i64,
    pub model: String,
    pub time_created: i64,
    pub time_updated: i64,
}

#[derive(Debug)]
pub struct UsageSessionData {
    pub session: RowUsageSession,
    pub model: Model,
}

#[derive(Debug, Deserialize)]
pub struct MessageCacheTokens {
    pub write: i64,
    pub read: i64,
}

#[derive(Debug, Deserialize)]
pub struct MessageTokens {
    // only in v1
    pub total: Option<i64>,
    pub input: i64,
    pub output: i64,
    pub reasoning: i64,
    pub cache: MessageCacheTokens,
}

#[derive(Debug, Deserialize)]
pub struct MessageData {
    // v2 only fields schema
    pub model: Option<Model>,
    // v1 only fields schema
    pub role: Option<String>,
    #[serde(rename = "providerID")]
    pub provider_id: Option<String>,
    #[serde(rename = "modelID")]
    pub model_id: Option<String>,
    // can be None in some cases (v1 schema)
    pub variant: Option<String>,
    // shared
    pub cost: Option<f64>,
    pub tokens: Option<MessageTokens>,
    pub agent: Option<String>,
    pub finish: Option<String>,
}

#[derive(Debug, sqlx::FromRow)]
pub struct RowSessionMessage {
    pub id: String,
    pub session_id: String,
    pub role: Option<String>,
    pub time_created: i64,
    pub time_updated: i64,
    pub data: String,
}

#[allow(dead_code)]
#[derive(Debug, Deserialize)]
pub struct SessionMessageItem {
    pub id: String,
    pub session_id: String,
    pub cost: f64,
    pub tokens_input: i64,
    pub tokens_output: i64,
    pub tokens_reasoning: i64,
    pub tokens_cache_read: i64,
    pub tokens_cache_write: i64,
    pub time_created: i64,
    pub time_updated: i64,
}

#[allow(dead_code)]
#[derive(Debug)]
pub struct SessionMessageData {
    pub message: SessionMessageItem,
    pub model: Model,
}

pub struct CollectClient {
    pool: Option<SqlitePool>,
}

#[derive(Deserialize, sqlx::FromRow)]
pub struct TableExists {
    pub exists: i8,
}

impl CollectClient {
    pub fn new() -> Self {
        CollectClient { pool: None }
    }

    pub fn get_db_path() -> String {
        let home = dirs::home_dir().unwrap();
        let db_path = home
            .join(".local")
            .join("share")
            .join("opencode")
            .join("opencode.db");
        db_path.to_str().unwrap().to_string()
    }

    async fn table_exists(&self, table_name: &str) -> anyhow::Result<bool> {
        let pool = match &self.pool {
            Some(pool) => pool,
            None => {
                anyhow::bail!("Database pool is not initialized. Call open_pool() first.");
            }
        };

        let exists: bool = sqlx::query_scalar(
            "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?)",
        )
        .bind(&table_name)
        .fetch_one(pool)
        .await?;

        Ok(exists)
    }

    pub async fn open_pool(&mut self) -> anyhow::Result<()> {
        let db_path = Self::get_db_path();
        let pool = SqlitePool::connect(&db_path).await?;
        self.pool = Some(pool);
        Ok(())
    }

    pub async fn get_usage_sessions(
        &self,
        min_updated_at: Option<i64>,
    ) -> anyhow::Result<Vec<UsageSessionData>> {
        let time_updated = match min_updated_at {
            Some(value) => value,
            None => 0, // Default to 0 if no value is provided
        };
        let pool = match &self.pool {
            Some(pool) => pool,
            None => {
                anyhow::bail!("Database pool is not initialized. Call open_pool() first.");
            }
        };

        let use_v2_table = self.table_exists("session_v2").await?;

        let sessions = if use_v2_table {
            sqlx::query_as::<_, RowUsageSession>(
                r#"
        SELECT
            id,
            project_id,
            slug,
            title,
            cost,
            tokens_input,
            tokens_output,
            tokens_reasoning,
            tokens_cache_read,
            tokens_cache_write,
            model,
            time_created,
            time_updated
        FROM session_v2
        WHERE time_updated > ?1
        ORDER BY time_created DESC
        "#,
            )
        } else {
            sqlx::query_as::<_, RowUsageSession>(
                r#"
        SELECT
            id,
            project_id,
            slug,
            title,
            cost,
            tokens_input,
            tokens_output,
            tokens_reasoning,
            tokens_cache_read,
            tokens_cache_write,
            model,
            time_created,
            time_updated
        FROM session
        WHERE time_updated > ?1
        ORDER BY time_created DESC
        "#,
            )
        }
        .bind(time_updated)
        .fetch_all(pool)
        .await?;

        let extended_sessions = sessions
            .into_iter()
            .map(|session| {
                let model_result = serde_json::from_str::<Model>(&session.model);
                let model = match model_result {
                    Ok(model) => model,
                    Err(e) => {
                        eprintln!(
                            "Failed to deserialize model for session {}: {}",
                            session.id, e
                        );
                        Model {
                            id: "unknown".to_string(),
                            provider_id: "unknown".to_string(),
                            variant: None,
                        }
                    }
                };

                UsageSessionData { session, model }
            })
            .collect();

        Ok(extended_sessions)
    }

    pub async fn get_session_messages(
        &self,
        min_updated_at: Option<i64>,
    ) -> anyhow::Result<Vec<SessionMessageData>> {
        let time_updated = match min_updated_at {
            Some(value) => value,
            None => 0, // Default to 0 if no value is provided
        };
        let pool = match &self.pool {
            Some(pool) => pool,
            None => {
                anyhow::bail!("Database pool is not initialized. Call open_pool() first.");
            }
        };

        let use_v2_table = self.table_exists("session_message").await?;
        let messages = if use_v2_table {
            sqlx::query_as::<_, RowSessionMessage>(
                r#"
        SELECT
            id,
            session_id,
            type as role,
            time_created,
            time_updated,
            data
        FROM session_message
        WHERE time_updated > ?1
        ORDER BY time_created DESC
        "#,
            )
        } else {
            sqlx::query_as::<_, RowSessionMessage>(
                r#"
        SELECT
            id,
            session_id,
            NULL as role,
            time_created,
            time_updated,
            data
        FROM message
        WHERE time_updated > ?1
        ORDER BY time_created DESC
        "#,
            )
        }
        .bind(time_updated)
        .fetch_all(pool)
        .await?;

        let extended_messages: Vec<SessionMessageData> = messages
            .into_iter()
            .filter_map(|message| {
                let message_result = serde_json::from_str::<MessageData>(&message.data);
                let message_data = match message_result {
                    Ok(message_data) => message_data,
                    Err(e) => {
                        eprintln!(
                            "Failed to deserialize data for message {}: {}",
                            message.id, e
                        );
                        return None;
                    }
                };

                let role = match message.role {
                    Some(role) => Some(role),
                    None => message_data.role,
                };
                if role.is_none() || role.unwrap() != "assistant" {
                    return None;
                }

                if message_data.finish.is_some() && message_data.finish.unwrap() == "error" {
                    return None;
                }

                let model = if let Some(model) = message_data.model {
                    model
                } else {
                    Model {
                        id: message_data.model_id.unwrap_or("unknown".to_owned()),
                        provider_id: message_data.provider_id.unwrap_or("unknown".to_owned()),
                        variant: message_data.variant,
                    }
                };
                if message_data.tokens.is_none() {
                    eprintln!("Tokens are missing for message {}", message.id);
                    return None;
                }

                let tokens = message_data.tokens.unwrap();

                let message_item = SessionMessageItem {
                    id: message.id,
                    session_id: message.session_id,
                    cost: message_data.cost.unwrap_or(0.0),
                    tokens_input: tokens.input,
                    tokens_output: tokens.output,
                    tokens_reasoning: tokens.reasoning,
                    tokens_cache_read: tokens.cache.read,
                    tokens_cache_write: tokens.cache.write,
                    time_created: message.time_created,
                    time_updated: message.time_updated,
                };

                Some(SessionMessageData {
                    message: message_item,
                    model,
                })
            })
            .collect();

        Ok(extended_messages)
    }
}