use std::collections::BTreeMap;
use std::io::{self, BufRead, BufReader};
use std::path::Path;
use std::process::{Command, Output};
use std::time::Instant;
use chrono::{DateTime, Utc};
use crate::request_log::{self, GhOutcome, LogRecord, RecordKind, Source};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Category {
Api,
Subcommand,
Local,
}
impl Category {
#[must_use]
pub fn from_command(command: &[String]) -> Self {
match command.first().map(String::as_str) {
Some("api") => Self::Api,
Some(first) if first.starts_with('-') => Self::Local,
_ => Self::Subcommand,
}
}
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Api => "api",
Self::Subcommand => "subcommand",
Self::Local => "local",
}
}
}
#[must_use]
pub fn source_str(source: Source) -> &'static str {
match source {
Source::Cli => "cli",
Source::Mcp => "mcp",
Source::Daemon => "daemon",
Source::Unknown => "unknown",
}
}
pub fn run_gh<I, S>(bin: &Path, args: I, label: &str, cwd: Option<&Path>) -> io::Result<Output>
where
I: IntoIterator<Item = S>,
S: AsRef<str>,
{
let argv: Vec<String> = args.into_iter().map(|s| s.as_ref().to_string()).collect();
let mut cmd = Command::new(bin);
cmd.args(&argv);
if let Some(dir) = cwd {
cmd.current_dir(dir);
}
let started = Instant::now();
let result = cmd.output();
let duration = started.elapsed();
let (exit_code, error) = match &result {
Ok(output) => (output.status.code(), None),
Err(e) => (None, Some(e.to_string())),
};
request_log::record_gh(GhOutcome {
label: label.to_string(),
argv,
exit_code,
duration,
error,
});
result
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct GhCounts {
pub by_category: BTreeMap<Category, u64>,
pub by_subcommand: BTreeMap<String, u64>,
pub by_source: BTreeMap<Source, u64>,
}
impl GhCounts {
#[must_use]
pub fn total(&self) -> u64 {
self.by_category.values().copied().sum()
}
#[must_use]
pub fn api_total(&self) -> u64 {
self.by_category
.iter()
.filter(|(cat, _)| **cat != Category::Local)
.map(|(_, n)| *n)
.sum()
}
#[must_use]
pub fn summary_line(&self) -> String {
let sources: Vec<String> = self
.by_source
.iter()
.map(|(s, n)| format!("{}={n}", source_str(*s)))
.collect();
let subs: Vec<String> = self
.by_subcommand
.iter()
.map(|(name, n)| format!("{name}={n}"))
.collect();
let local = self.by_category.get(&Category::Local).copied().unwrap_or(0);
format!(
"{} api call(s) (total {}, local {local}); by source: [{}]; by subcommand: [{}]",
self.api_total(),
self.total(),
sources.join(" "),
subs.join(" "),
)
}
#[must_use]
pub fn to_json(&self) -> serde_json::Value {
let by_category: BTreeMap<&str, u64> = self
.by_category
.iter()
.map(|(c, n)| (c.as_str(), *n))
.collect();
let by_source: BTreeMap<&str, u64> = self
.by_source
.iter()
.map(|(s, n)| (source_str(*s), *n))
.collect();
serde_json::json!({
"api_total": self.api_total(),
"total": self.total(),
"by_category": by_category,
"by_subcommand": self.by_subcommand,
"by_source": by_source,
})
}
fn tally(&mut self, command: &[String], source: Source) {
*self
.by_category
.entry(Category::from_command(command))
.or_default() += 1;
*self.by_subcommand.entry(command.join(" ")).or_default() += 1;
*self.by_source.entry(source).or_default() += 1;
}
}
#[must_use]
pub fn aggregate(
path: &Path,
since: Option<DateTime<Utc>>,
until: Option<DateTime<Utc>>,
source: Option<Source>,
) -> GhCounts {
let mut counts = GhCounts::default();
let Ok(file) = std::fs::File::open(path) else {
return counts; };
for line in BufReader::new(file).lines() {
let Ok(line) = line else { break };
if line.is_empty() {
continue;
}
let Ok(rec) = serde_json::from_str::<LogRecord>(&line) else {
continue; };
if rec.kind != RecordKind::Gh {
continue;
}
let rec_source = rec.source.unwrap_or(Source::Unknown);
if let Some(want) = source {
if rec_source != want {
continue;
}
}
if since.is_some() || until.is_some() {
let Some(ts) = parse_timestamp(&rec.timestamp) else {
continue; };
if since.is_some_and(|s| ts < s) || until.is_some_and(|u| ts > u) {
continue;
}
}
counts.tally(&rec.command, rec_source);
}
counts
}
pub(crate) fn parse_timestamp(ts: &str) -> Option<DateTime<Utc>> {
DateTime::parse_from_rfc3339(ts)
.ok()
.map(|dt| dt.with_timezone(&Utc))
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use std::io::Write;
fn cmd(parts: &[&str]) -> Vec<String> {
parts.iter().copied().map(String::from).collect()
}
#[test]
fn category_from_command_classifies_api_subcommand_and_local() {
assert_eq!(
Category::from_command(&cmd(&["api", "graphql"])),
Category::Api
);
assert_eq!(
Category::from_command(&cmd(&["api", "rate_limit"])),
Category::Api
);
assert_eq!(
Category::from_command(&cmd(&["pr", "list"])),
Category::Subcommand
);
assert_eq!(
Category::from_command(&cmd(&["repo", "view"])),
Category::Subcommand
);
assert_eq!(
Category::from_command(&cmd(&["--version"])),
Category::Local
);
assert_eq!(Category::from_command(&[]), Category::Subcommand);
}
const SAMPLE: &[&str] = &[
r#"{"kind":"gh","timestamp":"2026-07-21T10:00:00.000Z","command":["api","graphql"],"source":"daemon"}"#,
r#"{"kind":"gh","timestamp":"2026-07-21T10:01:00.000Z","command":["api","rate_limit"],"source":"daemon"}"#,
r#"{"kind":"gh","timestamp":"2026-07-21T10:02:00.000Z","command":["pr","list"],"source":"cli"}"#,
r#"{"kind":"gh","timestamp":"2026-07-21T10:03:00.000Z","command":["pr","list"],"source":"cli"}"#,
r#"{"kind":"gh","timestamp":"2026-07-21T10:04:00.000Z","command":["--version"],"source":"cli"}"#,
r#"{"kind":"http","timestamp":"2026-07-21T10:05:00.000Z","service":"jira"}"#,
"not json",
"",
];
fn write_log(lines: &[&str]) -> tempfile::NamedTempFile {
let mut f = tempfile::NamedTempFile::new().unwrap();
for line in lines {
writeln!(f, "{line}").unwrap();
}
f.flush().unwrap();
f
}
fn utc(ts: &str) -> DateTime<Utc> {
DateTime::parse_from_rfc3339(ts)
.unwrap()
.with_timezone(&Utc)
}
#[test]
fn aggregate_tallies_by_category_subcommand_and_source_ignoring_non_gh() {
let f = write_log(SAMPLE);
let counts = aggregate(f.path(), None, None, None);
assert_eq!(counts.total(), 5); assert_eq!(counts.api_total(), 4); assert_eq!(counts.by_category[&Category::Api], 2);
assert_eq!(counts.by_category[&Category::Subcommand], 2);
assert_eq!(counts.by_category[&Category::Local], 1);
assert_eq!(counts.by_subcommand["pr list"], 2);
assert_eq!(counts.by_subcommand["api graphql"], 1);
assert_eq!(counts.by_source[&Source::Daemon], 2);
assert_eq!(counts.by_source[&Source::Cli], 3);
}
#[test]
fn aggregate_filters_by_source() {
let f = write_log(SAMPLE);
let counts = aggregate(f.path(), None, None, Some(Source::Daemon));
assert_eq!(counts.total(), 2);
assert_eq!(counts.api_total(), 2);
assert!(!counts.by_source.contains_key(&Source::Cli));
}
#[test]
fn aggregate_filters_by_time_window() {
let f = write_log(SAMPLE);
let counts = aggregate(
f.path(),
Some(utc("2026-07-21T10:02:00.000Z")),
Some(utc("2026-07-21T10:03:30.000Z")),
None,
);
assert_eq!(counts.total(), 2);
assert_eq!(counts.by_subcommand["pr list"], 2);
}
#[test]
fn aggregate_missing_file_is_empty() {
let counts = aggregate(
std::path::Path::new("/nonexistent/omni-dev/log.jsonl"),
None,
None,
None,
);
assert_eq!(counts, GhCounts::default());
assert_eq!(counts.total(), 0);
}
#[test]
fn source_str_covers_every_source() {
assert_eq!(source_str(Source::Cli), "cli");
assert_eq!(source_str(Source::Mcp), "mcp");
assert_eq!(source_str(Source::Daemon), "daemon");
assert_eq!(source_str(Source::Unknown), "unknown");
}
#[test]
fn aggregate_drops_undateable_record_when_windowed() {
let lines = &[
r#"{"kind":"gh","timestamp":"not-a-timestamp","command":["pr","list"],"source":"cli"}"#,
r#"{"kind":"gh","timestamp":"2026-07-21T10:02:00.000Z","command":["pr","list"],"source":"cli"}"#,
];
let f = write_log(lines);
let windowed = aggregate(
f.path(),
Some(utc("2026-07-21T00:00:00.000Z")),
Some(utc("2026-07-22T00:00:00.000Z")),
None,
);
assert_eq!(windowed.total(), 1); assert_eq!(aggregate(f.path(), None, None, None).total(), 2);
}
#[test]
fn to_json_has_string_keys_and_totals() {
let f = write_log(SAMPLE);
let v = aggregate(f.path(), None, None, None).to_json();
assert_eq!(v["api_total"], 4);
assert_eq!(v["total"], 5);
assert_eq!(v["by_category"]["api"], 2);
assert_eq!(v["by_category"]["local"], 1);
assert_eq!(v["by_subcommand"]["pr list"], 2);
assert_eq!(v["by_source"]["daemon"], 2);
}
#[test]
fn summary_line_mentions_api_total() {
let f = write_log(SAMPLE);
let s = aggregate(f.path(), None, None, None).summary_line();
assert!(s.contains("4 api call"), "summary was: {s}");
}
}