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 {
pub total: Option<i64>,
pub input: i64,
pub output: i64,
pub reasoning: i64,
pub cache: MessageCacheTokens,
}
#[derive(Debug, Deserialize)]
pub struct MessageData {
pub model: Option<Model>,
pub role: Option<String>,
#[serde(rename = "providerID")]
pub provider_id: Option<String>,
#[serde(rename = "modelID")]
pub model_id: Option<String>,
pub variant: Option<String>,
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, };
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, };
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)
}
}