use serde_json::Value;
use std::{
collections::HashMap,
ffi::OsStr,
path::{Path, PathBuf},
};
use crate::{
apps::{normalize_model, value_hash},
error::AppError,
io::{load, progress::Progress},
model::{AppKind, UsageEntry},
};
const MAX_DEPTH: usize = 2;
const DSH_LINE_NEEDLES: [&str; 2] = ["\"assistant/message\"", "\"model/selection\""];
fn dsh_base(home: &Path, env: Option<&OsStr>) -> PathBuf {
load::env_abs_path("DSH_HOME", env, home).unwrap_or_else(|| home.join(".dsh"))
}
pub fn collect(threads: Option<usize>) -> Result<Vec<UsageEntry>, AppError> {
let base =
dsh_base(&load::home_dir()?, std::env::var_os("DSH_HOME").as_deref()).join("sessions");
if !base.is_dir() {
return Ok(Vec::new());
}
collect_from_with(&base, threads)
}
#[cfg(test)]
pub fn collect_from(base: &Path) -> Result<Vec<UsageEntry>, AppError> {
collect_from_with(base, None)
}
fn collect_from_with(base: &Path, threads: Option<usize>) -> Result<Vec<UsageEntry>, AppError> {
let files: Vec<PathBuf> = load::discover_files(base, "zstd", MAX_DEPTH)
.into_iter()
.filter(|f| load::file_name_str(f) == "session.jsonl.zstd")
.collect();
let threads = threads.unwrap_or_else(|| load::auto_threads(files.len()));
let progress = Progress::start("dsh", load::total_bytes(&files), files.len());
let per_file: Vec<HashMap<String, UsageEntry>> =
load::map_files(&files, threads, &progress, parse_session);
progress.finish();
let mut candidates: HashMap<String, UsageEntry> = HashMap::new();
for map in per_file {
for (key, entry) in map {
candidates.entry(key).or_insert(entry);
}
}
Ok(candidates.into_values().collect())
}
fn parse_session(file: &Path, progress: &Progress) -> HashMap<String, UsageEntry> {
let mut candidates: HashMap<String, UsageEntry> = HashMap::new();
let mut first_seen = false;
let mut session_id = "unknown".to_string();
let mut header_ts: Option<i64> = None;
let mut current_model: Option<String> = None;
if let Err(e) = load::for_each_jsonl_zstd_progress(file, &DSH_LINE_NEEDLES, progress, |entry| {
if !first_seen {
first_seen = true;
if load::str_get(&entry, &["type"]) != Some("session") {
return false;
}
session_id = load::str_get(&entry, &["id"])
.unwrap_or("unknown")
.to_string();
header_ts = entry.get("createdAt").and_then(load::timestamp_to_epoch);
return true;
}
match load::str_get(&entry, &["type"]) {
Some("model/selection") => {
if let Some(m) = load::str_get(&entry, &["data", "model"]) {
current_model = Some(normalize_model(m));
}
}
Some("assistant/message") => {
if let Some((key, record)) =
parse_entry(&entry, &session_id, header_ts, current_model.as_deref())
{
candidates.entry(key).or_insert(record);
}
}
_ => {}
}
true
}) {
load::warn_file(&e);
progress.note_error();
}
candidates
}
fn parse_entry(
entry: &Value,
session_id: &str,
header_ts: Option<i64>,
current_model: Option<&str>,
) -> Option<(String, UsageEntry)> {
let data = entry.get("data").filter(|d| d.is_object())?;
let usage = data.get("usage").filter(|u| u.is_object())?;
let input = load::u64_get(usage, &["inputTokens"]);
let output = load::u64_get(usage, &["outputTokens"]);
let cache_read = load::u64_get(usage, &["cacheReadTokens"]);
let cache_write = load::u64_get(usage, &["cacheWriteTokens"]);
if input == 0 && output == 0 && cache_read == 0 && cache_write == 0 {
return None;
}
let model = load::str_get(data, &["message", "source", "model"])
.or(current_model)
.map_or_else(|| "unknown".to_string(), normalize_model);
let created_at = entry
.get("time")
.and_then(load::timestamp_to_epoch)
.or(header_ts)
.unwrap_or_else(load::now_epoch);
let key = match load::str_get(data, &["message", "id"]).filter(|s| !s.is_empty()) {
Some(id) => format!("id:{id}"),
None => format!("hash:{}", value_hash(entry)),
};
Some((
key,
UsageEntry::new(
AppKind::Dsh,
model,
Some(session_id.to_string()),
created_at,
input,
output,
cache_read,
cache_write,
None,
),
))
}
#[cfg(test)]
#[path = "tests/dsh_test.rs"]
mod tests;