diff --git a/.gitignore b/.gitignore index 8255403..6fdae7b 100644 --- a/.gitignore +++ b/.gitignore @@ -36,4 +36,5 @@ node_modules/ DEPLOY.md PUBLISH.md tmp-statsai-*.sqlite* +/statsai.sqlite* .wrangler/ diff --git a/Cargo.lock b/Cargo.lock index ade35e3..3899396 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -437,6 +437,27 @@ dependencies = [ "typenum", ] +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + [[package]] name = "dbus" version = "0.9.11" @@ -2170,6 +2191,12 @@ dependencies = [ "twox-hash", ] +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "same-file" version = "1.0.6" @@ -2363,6 +2390,7 @@ dependencies = [ "anyhow", "chrono", "clap", + "csv", "getrandom 0.2.17", "hex", "hmac", diff --git a/Cargo.toml b/Cargo.toml index a4cb973..7364866 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,7 @@ anyhow = "1" base64 = "0.22" chrono = { version = "0.4", features = ["serde"] } clap = { version = "4", features = ["derive"] } +csv = "1" hex = "0.4" libc = "0.2" rusqlite = { version = "0.32", features = ["bundled"] } diff --git a/crates/statsai-adapters/src/archive.rs b/crates/statsai-adapters/src/archive.rs index 1731bc9..48a8a89 100644 --- a/crates/statsai-adapters/src/archive.rs +++ b/crates/statsai-adapters/src/archive.rs @@ -1,9 +1,10 @@ mod mutations; use super::{ - canonical_display, codex_project_context_from_value, codex_usage_roots, collect_jsonl_files, - expand_home_path, grok_sessions_root, model_from_nested_value, open_sqlite_readonly, - read_bounded_jsonl_line, resolve_project_context, source_root_path, + canonical_display, codex_project_context_from_value, codex_quota_observation, + codex_usage_counts_from_value, codex_usage_roots, collect_jsonl_files, expand_home_path, + file_modified_timestamp, grok_sessions_root, model_from_nested_value, open_sqlite_readonly, + read_bounded_jsonl_line, resolve_project_context, source_root_path, subtract_usage_counts, timestamp_from_nested_value, BoundedLineRead, ProjectContextCache, CLAUDE_CODE_PROVIDER, CODEX_PROVIDER, GROK_BUILD_PROVIDER, MAX_JSONL_RECORD_BYTES, OPENCODE_PROVIDER, }; @@ -21,8 +22,8 @@ use statsai_core::{ archive_artifact_metadata_signature, archive_content_id, archive_conversation_id, archive_item_id, hash_text, ArchiveArtifactDependency, ArchiveCompleteness, ArchiveContentKind, ArchiveContentPart, ArchiveConversation, ArchiveItem, ArchiveItemKind, ArchiveRole, - CoverageStatus, ModelInfo, ProjectInfo, SourceLocation, TraceEdit, UsageCounts, - ARCHIVE_CONVERSATION_SCHEMA_VERSION, + CoverageStatus, ModelInfo, ProjectInfo, QuotaObservationRecordV1, SourceLocation, TraceEdit, + UsageCounts, ARCHIVE_CONVERSATION_SCHEMA_VERSION, }; use std::collections::{BTreeMap, HashMap, HashSet}; use std::fs::File; @@ -57,6 +58,7 @@ pub struct ArchiveScan { pub conversations: Vec, pub artifact_dependencies: Vec, pub trace_edits: Vec, + pub quota_observations: Vec, pub trace_coverage: CoverageStatus, pub diagnostics: ArchiveScanDiagnostics, } @@ -67,6 +69,7 @@ impl Default for ArchiveScan { conversations: Vec::new(), artifact_dependencies: Vec::new(), trace_edits: Vec::new(), + quota_observations: Vec::new(), // A scan that reconstructed nothing has measured nothing. Only a // provider whose mutations this module actually parses may claim // complete trace coverage, so that a new adapter cannot silently @@ -210,12 +213,70 @@ fn collect_codex( &mut scan.trace_edits, &mut scan.trace_coverage, )?); + scan.quota_observations + .extend(collect_codex_quota_observations(source, &path)?); } scan.artifact_dependencies = finish_artifact_dependencies(artifact_dependencies); finish_diagnostics(&mut scan); Ok(scan) } +fn collect_codex_quota_observations( + source: &SourceLocation, + path: &Path, +) -> Result> { + let file = File::open(path).with_context(|| format!("read {}", path.display()))?; + let mut reader = BufReader::new(file); + let fallback_timestamp = file_modified_timestamp(path).unwrap_or_else(Utc::now); + let mut previous_totals: Option = None; + let mut observations = Vec::new(); + let mut line_bytes = Vec::new(); + let mut line_number = 0usize; + loop { + let status = read_bounded_jsonl_line(&mut reader, &mut line_bytes, MAX_JSONL_RECORD_BYTES)?; + if status == BoundedLineRead::Eof { + break; + } + line_number = line_number.saturating_add(1); + if status == BoundedLineRead::Oversized { + continue; + } + let Ok(line) = std::str::from_utf8(&line_bytes) else { + continue; + }; + let Ok(value) = serde_json::from_str::(line) else { + continue; + }; + if value.get("type").and_then(Value::as_str) != Some("event_msg") + || value.pointer("/payload/type").and_then(Value::as_str) != Some("token_count") + { + continue; + } + let info = value.pointer("/payload/info"); + let total_usage = info + .and_then(|info| info.get("total_token_usage")) + .map(codex_usage_counts_from_value); + let usage_sample = info + .and_then(|info| info.get("last_token_usage")) + .map(codex_usage_counts_from_value) + .or_else(|| { + total_usage + .as_ref() + .map(|total| subtract_usage_counts(total, previous_totals.as_ref())) + }); + if let Some(total_usage) = total_usage { + previous_totals = Some(total_usage); + } + let observed_at = timestamp_from_nested_value(&value).unwrap_or(fallback_timestamp); + if let Some(observation) = + codex_quota_observation(source, path, line_number, observed_at, usage_sample, &value) + { + observations.push(observation); + } + } + Ok(observations) +} + fn collect_codex_file( source: &SourceLocation, path: &Path, @@ -1268,6 +1329,7 @@ fn collect_opencode( .collect(), artifact_dependencies: finish_artifact_dependencies(artifact_dependencies), trace_edits: Vec::new(), + quota_observations: Vec::new(), trace_coverage: CoverageStatus::Unavailable, diagnostics: ArchiveScanDiagnostics { files_scanned: 1, diff --git a/crates/statsai-adapters/src/lib.rs b/crates/statsai-adapters/src/lib.rs index 50dadf7..0a071ab 100644 --- a/crates/statsai-adapters/src/lib.rs +++ b/crates/statsai-adapters/src/lib.rs @@ -12,10 +12,12 @@ use statsai_core::{ task_title_from_prompt, task_title_is_generic, task_title_is_weak_signal, task_title_signal_score, title_topic_tokens, BillingPeriod, Confidence, CostAccumulator, CostInfo, EventId, EventSource, IdentitySource, LatencySource, LocationOrigin, MetricStats, - ModelInfo, ParseEvidence, PrivacyInfo, PrivacyMode, ProjectInfo, ReasoningLevel, RuntimeInfo, - SessionInfo, SourceKind, SourceLocation, SubscriptionStatus, SummaryMetadata, SummaryMetrics, - TaskSpan, UsageCounts, UsageEvent, UsageSummary, TASK_SPAN_SCHEMA_VERSION, - USAGE_EVENT_SCHEMA_VERSION, USAGE_SUMMARY_SCHEMA_VERSION, + ModelInfo, ParseEvidence, PrivacyInfo, PrivacyMode, ProjectInfo, QuotaCreditsV1, + QuotaObservationRecordV1, QuotaObservationV1, QuotaStatusV1, QuotaUsageLinkKind, + QuotaWindowObservationV1, ReasoningLevel, RuntimeInfo, SessionInfo, SourceKind, SourceLocation, + SubscriptionStatus, SummaryMetadata, SummaryMetrics, TaskSpan, UsageCounts, UsageEvent, + UsageSummary, QUOTA_OBSERVATION_SCHEMA_VERSION, QUOTA_WINDOW_OBSERVATION_SCHEMA_VERSION, + TASK_SPAN_SCHEMA_VERSION, USAGE_EVENT_SCHEMA_VERSION, USAGE_SUMMARY_SCHEMA_VERSION, }; use statsai_pricing::{ estimate_cost_at, normalize_model_name, pricing_changes_between, unknown_cost, @@ -39,7 +41,7 @@ const PROVIDER_RECORD_EVENT_KEY_VERSION: &str = "provider_record_usage_event.v1" const SCAN_CACHE_SIGNATURE_VERSION: &str = "scan-cache.v1"; // Invalidate unchanged-file scan cache entries whenever provider parsing semantics change, // so historical sessions get rescanned for runtime, pricing, and project context updates. -const CODEX_SCAN_CACHE_PARSER_REVISION: &str = "task-spans.v26"; +const CODEX_SCAN_CACHE_PARSER_REVISION: &str = "quota-history.v27"; const CLAUDE_SCAN_CACHE_PARSER_REVISION: &str = "task-spans.v23"; const OPENCODE_SCAN_CACHE_PARSER_REVISION: &str = "task-spans.v15"; const GROK_BUILD_SCAN_CACHE_PARSER_REVISION: &str = "task-spans.v20"; @@ -196,6 +198,7 @@ pub struct AdapterScan { pub events: Vec, pub summaries: Vec, pub task_spans: Vec, + pub quota_observations: Vec, pub diagnostics: ScanDiagnostics, pub verified_source_state: Option, } @@ -2593,6 +2596,7 @@ fn parse_codex_file( let mut current_thread_id: Option = None; let session_raw = codex_session_id(usage_root, path); let mut records = Vec::new(); + let mut quota_observation_indices = HashMap::new(); let mut project_cache = ProjectContextCache::new(); let mut line_bytes = Vec::new(); let mut index = 0usize; @@ -2619,7 +2623,7 @@ fn parse_codex_file( } ctx.scan.diagnostics.raw_rows += 1; let line_kind = codex_line_kind(line); - if line_kind == CodexLineKind::Irrelevant { + if line_kind == CodexLineKind::Irrelevant && !is_codex_quota_line_structurally(line) { continue; } if line_kind == CodexLineKind::ResponseItemMessage { @@ -2918,6 +2922,7 @@ fn parse_codex_file( } else { codex_headless_usage_value(&value).map(codex_usage_counts_from_value) }; + let quota_usage_sample = usage.clone(); let (timestamp, timestamp_inferred) = timestamp_from_nested_value(&value) .map(|timestamp| (timestamp, false)) @@ -2925,6 +2930,17 @@ fn parse_codex_file( if timestamp_inferred { ctx.scan.diagnostics.timestamp_fallbacks += 1; } + if let Some(quota) = codex_quota_observation( + ctx.source, + path, + index, + timestamp, + quota_usage_sample, + &value, + ) { + quota_observation_indices.insert(index, ctx.scan.quota_observations.len()); + ctx.scan.quota_observations.push(quota); + } let explicit_model = with_reasoning_state(codex_model_from_value(&value, None), ¤t_reasoning); @@ -3131,8 +3147,8 @@ fn parse_codex_file( let Some(usage) = usage else { continue; }; - for line_number in turn.usage_lines { - consumed_usage_lines.insert(line_number); + for line_number in &turn.usage_lines { + consumed_usage_lines.insert(*line_number); } if record.usage.is_some() { consumed_usage_lines.insert(record.line_number); @@ -3184,6 +3200,17 @@ fn parse_codex_file( dedupe_salt: None, }, ); + let mut linked_quota_lines = turn.usage_lines.clone(); + if record.usage.is_some() { + linked_quota_lines.push(record.line_number); + } + link_quota_observations( + ctx.scan, + "a_observation_indices, + &linked_quota_lines, + &event.event_id, + QuotaUsageLinkKind::TurnEvent, + ); let task_span = if ctx.options.should_collect_tasks() { let event_id = event.event_id.clone(); let event_cost = event.cost.estimated_api_equivalent_usd; @@ -3402,12 +3429,289 @@ fn parse_codex_file( dedupe_salt: None, }, ); + link_quota_observations( + ctx.scan, + "a_observation_indices, + &[record.line_number], + &event.event_id, + QuotaUsageLinkKind::RecordEvent, + ); push_deduped(ctx.scan, ctx.seen, event); } Ok(()) } +fn link_quota_observations( + scan: &mut AdapterScan, + quota_observation_indices: &HashMap, + line_numbers: &[usize], + event_id: &EventId, + link_kind: QuotaUsageLinkKind, +) { + if line_numbers.is_empty() { + return; + } + for line_number in line_numbers { + let Some(record) = quota_observation_indices + .get(line_number) + .and_then(|index| scan.quota_observations.get_mut(*index)) + else { + continue; + }; + if record + .observation + .usage_sample + .as_ref() + .is_some_and(|usage| usage.computed_total() > 0) + { + record.observation.usage_event_id = Some(event_id.clone()); + record.observation.usage_link_kind = link_kind; + } + } +} + +pub(crate) fn codex_quota_observation( + source: &SourceLocation, + path: &Path, + line_number: usize, + observed_at: DateTime, + usage_sample: Option, + value: &Value, +) -> Option { + if !is_codex_token_count(value) { + return None; + } + let rate_limits = value.pointer("/payload/rate_limits")?.as_object()?; + let raw_rate_limits = Value::Object(rate_limits.clone()); + let raw_json = serde_json::to_string(&raw_rate_limits).ok()?; + let payload_hash = hash_text(&raw_json); + let source_file_path_hash = hash_text(&canonical_display(path)); + let source_record_id = format!("quota:{source_file_path_hash}:{line_number}"); + let observation_id = format!( + "quota_observation_{}", + &hash_text(&format!("{}:{source_record_id}", source.source_id.0))[..32] + ); + let semantic_fingerprint = hash_text(&format!( + "quota_semantic.v1:{}:{}:{payload_hash}", + source.provider, + observed_at.to_rfc3339() + )); + let global_limit_id = quota_string_at_any( + &raw_rate_limits, + &["limit_id", "limitId", "limit_name", "limitName", "id"], + ); + let mut windows = Vec::new(); + for (slot, candidate) in rate_limits { + let Some(candidate) = candidate.as_object() else { + continue; + }; + let Some(window_minutes) = quota_u64_at_any( + candidate, + &["window_minutes", "windowMinutes", "duration_minutes"], + ) else { + continue; + }; + let Some(used_percent) = quota_f64_at_any( + candidate, + &["used_percent", "usedPercent", "percentage", "percent_used"], + ) else { + continue; + }; + let Some(resets_at_epoch_seconds) = + quota_i64_at_any(candidate, &["resets_at", "resetsAt", "reset_at", "resetAt"]) + else { + continue; + }; + let Some(resets_at) = Utc.timestamp_opt(resets_at_epoch_seconds, 0).single() else { + continue; + }; + let limit_id = quota_string_at_any_from_map( + candidate, + &["limit_id", "limitId", "limit_name", "limitName", "id"], + ) + .or_else(|| global_limit_id.clone()); + let window_observation_id = format!( + "quota_window_observation_{}", + &hash_text(&format!("{observation_id}:{slot}"))[..32] + ); + windows.push(QuotaWindowObservationV1 { + schema_version: QUOTA_WINDOW_OBSERVATION_SCHEMA_VERSION.to_string(), + window_observation_id, + observation_id: observation_id.clone(), + provider_slot: slot.clone(), + limit_id, + window_minutes, + used_percent, + resets_at, + resets_at_epoch_seconds, + }); + } + + let credits_value = rate_limits.get("credits").and_then(Value::as_object); + let balance_raw = credits_value + .and_then(|credits| credits.get("balance")) + .cloned(); + let balance = balance_raw.as_ref().and_then(normalize_quota_decimal); + let status = QuotaStatusV1 { + plan_type: quota_string_at_any( + &raw_rate_limits, + &["plan_type", "planType", "plan", "subscription_type"], + ), + individual_limit: rate_limits + .get("individual_limit") + .or_else(|| rate_limits.get("individualLimit")) + .cloned(), + spend_control_state: quota_string_at_any( + &raw_rate_limits, + &["spend_control_state", "spendControlState"], + ) + .or_else(|| { + raw_rate_limits + .pointer("/spend_control/state") + .and_then(Value::as_str) + .map(ToOwned::to_owned) + }), + reached_type: quota_string_at_any( + &raw_rate_limits, + &["reached_type", "reachedType", "limit_reached_type"], + ), + credits: QuotaCreditsV1 { + has_credits: credits_value + .and_then(|credits| credits.get("has_credits")) + .and_then(Value::as_bool), + unlimited: credits_value + .and_then(|credits| credits.get("unlimited")) + .and_then(Value::as_bool), + balance, + balance_raw, + }, + }; + Some(QuotaObservationRecordV1 { + observation: QuotaObservationV1 { + schema_version: QUOTA_OBSERVATION_SCHEMA_VERSION.to_string(), + observation_id, + semantic_fingerprint, + provider: source.provider.clone(), + source_id: source.source_id.clone(), + provider_account_id: None, + observed_at, + source_file_path_hash, + source_record_id, + source_line_number: line_number as u64, + payload_hash, + usage_sample, + usage_event_id: None, + usage_link_kind: QuotaUsageLinkKind::None, + status, + }, + windows, + raw_rate_limits, + }) +} + +fn quota_string_at_any(value: &Value, keys: &[&str]) -> Option { + value + .as_object() + .and_then(|map| quota_string_at_any_from_map(map, keys)) +} + +fn quota_string_at_any_from_map( + map: &serde_json::Map, + keys: &[&str], +) -> Option { + keys.iter() + .filter_map(|key| map.get(*key)) + .find_map(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn quota_u64_at_any(map: &serde_json::Map, keys: &[&str]) -> Option { + keys.iter() + .filter_map(|key| map.get(*key)) + .find_map(Value::as_u64) +} + +fn quota_i64_at_any(map: &serde_json::Map, keys: &[&str]) -> Option { + keys.iter() + .filter_map(|key| map.get(*key)) + .find_map(Value::as_i64) +} + +fn quota_f64_at_any(map: &serde_json::Map, keys: &[&str]) -> Option { + keys.iter() + .filter_map(|key| map.get(*key)) + .find_map(Value::as_f64) + .filter(|value| value.is_finite()) +} + +fn normalize_quota_decimal(value: &Value) -> Option { + let text: Cow<'_, str> = match value { + Value::String(value) => Cow::Borrowed(value.trim()), + Value::Number(value) => Cow::Owned(value.to_string()), + Value::Null => return None, + _ => return None, + }; + if text.is_empty() || text.len() > 4_096 { + return None; + } + let (negative, unsigned) = text + .strip_prefix('-') + .map_or((false, text.as_ref()), |value| (true, value)); + let (mantissa, exponent) = + unsigned + .split_once(['e', 'E']) + .map_or((unsigned, 0i32), |(mantissa, exponent)| { + exponent + .parse::() + .ok() + .filter(|value| value.unsigned_abs() <= 4_096) + .map(|exponent| (mantissa, exponent)) + .unwrap_or(("", 0)) + }); + if mantissa.is_empty() { + return None; + } + let (whole, fraction) = mantissa.split_once('.').unwrap_or((mantissa, "")); + if whole.is_empty() && fraction.is_empty() + || !whole.bytes().all(|byte| byte.is_ascii_digit()) + || !fraction.bytes().all(|byte| byte.is_ascii_digit()) + { + return None; + } + let digits = format!("{whole}{fraction}"); + let decimal_position = i64::try_from(whole.len()).ok()? + i64::from(exponent); + let expanded = if decimal_position <= 0 { + let zeroes = usize::try_from(-decimal_position).ok()?; + format!("0.{}{}", "0".repeat(zeroes), digits) + } else if decimal_position >= i64::try_from(digits.len()).ok()? { + let zeroes = usize::try_from(decimal_position) + .ok()? + .saturating_sub(digits.len()); + format!("{digits}{}", "0".repeat(zeroes)) + } else { + let position = usize::try_from(decimal_position).ok()?; + format!("{}.{}", &digits[..position], &digits[position..]) + }; + let (whole, fraction) = expanded + .split_once('.') + .map_or((expanded.as_str(), ""), |parts| parts); + let whole = whole.trim_start_matches('0'); + let fraction = fraction.trim_end_matches('0'); + let whole = if whole.is_empty() { "0" } else { whole }; + let mut normalized = if fraction.is_empty() { + whole.to_string() + } else { + format!("{whole}.{fraction}") + }; + if negative && normalized != "0" { + normalized.insert(0, '-'); + } + Some(normalized) +} + #[derive(Debug, Clone)] struct CodexLineRecord { line_number: usize, @@ -3542,6 +3846,9 @@ fn merge_adapter_scan( } target.summaries.append(&mut source.summaries); target.task_spans.append(&mut source.task_spans); + target + .quota_observations + .append(&mut source.quota_observations); target.diagnostics.files_scanned = target .diagnostics .files_scanned @@ -5572,6 +5879,38 @@ enum CodexLineKind { HeadlessUsage, } +#[derive(Deserialize)] +struct CodexQuotaLineProbe { + #[serde(rename = "type")] + line_type: Option, + payload: Option, +} + +#[derive(Deserialize)] +struct CodexQuotaPayloadProbe { + #[serde(rename = "type")] + payload_type: Option, + rate_limits: Option, +} + +fn is_codex_quota_line_structurally(line: &str) -> bool { + if !line.contains("\"event_msg\"") + || !line.contains("\"token_count\"") + || !line.contains("\"rate_limits\"") + { + return false; + } + serde_json::from_str::(line) + .ok() + .is_some_and(|probe| { + probe.line_type.as_deref() == Some("event_msg") + && probe.payload.is_some_and(|payload| { + payload.payload_type.as_deref() == Some("token_count") + && payload.rate_limits.is_some() + }) + }) +} + fn codex_line_header(line: &str) -> &str { codex_prefix_at_char_boundary(line, 256) } @@ -7646,6 +7985,233 @@ mod tests { assert_eq!(source.path_label.as_deref(), Some("/tmp/codex-home")); } + #[test] + fn codex_quota_parser_is_anchored_and_preserves_modern_status_fields() { + let adapter = CodexAdapter; + let source = codex_source_for_root( + &adapter, + Path::new("/tmp/codex-home"), + LocationOrigin::Configured, + ); + let observed_at = Utc.with_ymd_and_hms(2026, 8, 20, 12, 0, 0).unwrap(); + let value = serde_json::json!({ + "timestamp": observed_at, + "type": "event_msg", + "payload": { + "type": "token_count", + "info": {"last_token_usage": {"total_tokens": 0}}, + "rate_limits": { + "limit_id": "codex_subscription", + "plan_type": "pro", + "individual_limit": null, + "spend_control_state": "allowed", + "reached_type": "weekly", + "primary": { + "used_percent": 12.5, + "window_minutes": 10080, + "resets_at": 1787832000 + }, + "credits": { + "has_credits": true, + "unlimited": false, + "balance": "0012.5000" + } + } + } + }); + let record = codex_quota_observation( + &source, + Path::new("/tmp/codex-home/sessions/thread.jsonl"), + 7, + observed_at, + Some(UsageCounts::default()), + &value, + ) + .expect("quota observation"); + + assert_eq!(record.windows.len(), 1); + assert_eq!(record.windows[0].provider_slot, "primary"); + assert_eq!(record.windows[0].window_minutes, 10_080); + assert_eq!( + record.windows[0].limit_id.as_deref(), + Some("codex_subscription") + ); + assert_eq!(record.observation.status.plan_type.as_deref(), Some("pro")); + assert_eq!( + record.observation.status.credits.balance.as_deref(), + Some("12.5") + ); + assert_eq!( + record.observation.status.credits.balance_raw, + Some(Value::String("0012.5000".to_string())) + ); + assert_eq!(record.observation.usage_link_kind, QuotaUsageLinkKind::None); + + let nested_as_text = serde_json::json!({ + "type": "event_msg", + "payload": { + "type": "user_message", + "message": value.to_string() + } + }); + assert!(codex_quota_observation( + &source, + Path::new("/tmp/thread.jsonl"), + 1, + observed_at, + None, + &nested_as_text, + ) + .is_none()); + } + + #[test] + fn codex_quota_parser_requires_integer_reset_epochs_and_leniently_reads_balances() { + let adapter = CodexAdapter; + let source = codex_source_for_root( + &adapter, + Path::new("/tmp/codex-home"), + LocationOrigin::Configured, + ); + let observed_at = Utc.with_ymd_and_hms(2026, 8, 20, 12, 0, 0).unwrap(); + for (balance, normalized) in [ + (serde_json::json!(14.25), Some("14.25")), + (serde_json::json!("1.25e-3"), Some("0.00125")), + (Value::Null, None), + ] { + let value = serde_json::json!({ + "type": "event_msg", + "payload": { + "type": "token_count", + "rate_limits": { + "primary": { + "used_percent": 1, + "window_minutes": 300, + "resets_at": "1787832000" + }, + "credits": {"balance": balance} + } + } + }); + let record = codex_quota_observation( + &source, + Path::new("/tmp/thread.jsonl"), + 1, + observed_at, + None, + &value, + ) + .expect("structural quota payload"); + assert!(record.windows.is_empty(), "string epochs are invalid"); + assert_eq!( + record.observation.status.credits.balance.as_deref(), + normalized + ); + } + assert!(codex_quota_observation( + &source, + Path::new("/tmp/thread.jsonl"), + 1, + observed_at, + None, + &serde_json::json!({ + "type": "event_msg", + "payload": {"type": "token_count", "rate_limits": "malformed"} + }), + ) + .is_none()); + } + + #[test] + fn codex_quota_links_consumed_samples_to_turn_events_and_preserves_zero_samples() { + let directory = tempfile::tempdir().expect("tempdir"); + let root = directory.path().join("codex"); + let sessions = root.join("sessions"); + std::fs::create_dir_all(&sessions).expect("sessions"); + let path = sessions.join("thread.jsonl"); + let mut fixture = File::create(&path).expect("fixture"); + for value in [ + serde_json::json!({ + "timestamp": "2026-08-20T12:00:00Z", + "type": "event_msg", + "payload": {"type": "task_started"} + }), + serde_json::json!({ + "timestamp": "2026-08-20T12:00:01Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": {"input_tokens": 10, "output_tokens": 5}, + "total_token_usage": {"input_tokens": 10, "output_tokens": 5} + }, + "rate_limits": { + "primary": { + "used_percent": 10, + "window_minutes": 10080, + "resets_at": 1787832000 + } + } + } + }), + serde_json::json!({ + "timestamp": "2026-08-20T12:00:02Z", + "type": "event_msg", + "payload": {"type": "task_complete"} + }), + serde_json::json!({ + "timestamp": "2026-08-20T12:00:03Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": {"total_tokens": 0}, + "total_token_usage": {"input_tokens": 10, "output_tokens": 5} + }, + "rate_limits": { + "primary": { + "used_percent": 11, + "window_minutes": 10080, + "resets_at": 1787832000 + } + } + } + }), + ] { + writeln!(fixture, "{value}").expect("write fixture"); + } + drop(fixture); + let source = codex_source_for_root(&CodexAdapter, &root, LocationOrigin::Configured); + let scan = scan_codex_source(&CodexAdapter, &source, &options()).expect("scan"); + + assert_eq!(scan.events.len(), 1); + assert_eq!(scan.quota_observations.len(), 2); + assert_eq!( + scan.quota_observations[0].observation.usage_link_kind, + QuotaUsageLinkKind::TurnEvent + ); + assert_eq!( + scan.quota_observations[0].observation.usage_event_id, + Some(scan.events[0].event_id.clone()) + ); + assert_eq!( + scan.quota_observations[1] + .observation + .usage_sample + .as_ref() + .map(UsageCounts::computed_total), + Some(0) + ); + assert_eq!( + scan.quota_observations[1].observation.usage_link_kind, + QuotaUsageLinkKind::None + ); + assert!(scan.quota_observations[1] + .observation + .usage_event_id + .is_none()); + } + #[test] fn claude_normalizes_projects_path_to_config_root() { let adapter = ClaudeCodeAdapter; @@ -9214,6 +9780,14 @@ mod tests { codex_line_kind(user_message), CodexLineKind::ResponseItemMessage ); + + let unrelated_event = r#"{"timestamp":"2026-06-03T09:36:26.000Z","type":"event_msg","payload":{"type":"agent_message","model":"gpt-incorrect"}}"#; + assert_eq!(codex_line_kind(unrelated_event), CodexLineKind::Irrelevant); + assert!(!is_codex_quota_line_structurally(reasoning)); + + let reordered_quota = r#"{"payload": {"rate_limits": {"primary": {"used_percent": 1, "window_minutes": 300, "resets_at": 1787832000}}, "type": "token_count"}, "type": "event_msg"}"#; + assert_eq!(codex_line_kind(reordered_quota), CodexLineKind::Irrelevant); + assert!(is_codex_quota_line_structurally(reordered_quota)); } #[test] @@ -12021,6 +12595,47 @@ mod tests { ); } + #[test] + fn codex_non_usage_event_messages_do_not_override_turn_model() { + let dir = tempfile::tempdir().expect("tempdir"); + let sessions = dir.path().join("sessions"); + std::fs::create_dir_all(&sessions).expect("sessions"); + let mut file = File::create(sessions.join("model.jsonl")).expect("fixture"); + file.write_all( + br#"{"timestamp":"2026-05-01T00:00:00Z","type":"turn_context","payload":{"model":"gpt-5.4"}} +"#, + ) + .expect("write context"); + file.write_all( + br#"{"timestamp":"2026-05-01T00:00:01Z","type":"event_msg","payload":{"type":"agent_message","model":"gpt-incorrect"}} +"#, + ) + .expect("write unrelated event"); + file.write_all( + br#"{"timestamp":"2026-05-01T00:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":80,"output_tokens":40,"total_tokens":120},"total_token_usage":{"input_tokens":80,"output_tokens":40,"total_tokens":120}}}} +"#, + ) + .expect("write usage"); + let source = SourceLocation::local_adapter( + CODEX_PROVIDER, + "test", + "0", + dir.path(), + LocationOrigin::Configured, + ); + + let scan = scan_codex_source(&CodexAdapter, &source, &options()).expect("scan"); + + assert_eq!(scan.events.len(), 1); + assert_eq!( + scan.events[0] + .model + .as_ref() + .and_then(|model| model.normalized_name.as_deref()), + Some("gpt-5.4") + ); + } + #[test] fn claude_adapter_does_not_infer_reasoning_level_from_thinking_model_name() { let dir = tempfile::tempdir().expect("tempdir"); diff --git a/crates/statsai-core/Cargo.toml b/crates/statsai-core/Cargo.toml index 92a4c03..8005e37 100644 --- a/crates/statsai-core/Cargo.toml +++ b/crates/statsai-core/Cargo.toml @@ -13,10 +13,10 @@ chrono.workspace = true hex.workspace = true schemars.workspace = true serde.workspace = true +serde_json.workspace = true sha2.workspace = true thiserror.workspace = true unicode-normalization.workspace = true [dev-dependencies] -serde_json.workspace = true tempfile.workspace = true diff --git a/crates/statsai-core/src/lib.rs b/crates/statsai-core/src/lib.rs index c5322a7..27789ff 100644 --- a/crates/statsai-core/src/lib.rs +++ b/crates/statsai-core/src/lib.rs @@ -2,6 +2,7 @@ mod archive; mod code_changes; +mod quota; mod tasks; use chrono::{DateTime, Utc}; @@ -12,6 +13,7 @@ use std::path::{Path, PathBuf}; pub use archive::*; pub use code_changes::*; +pub use quota::*; pub use tasks::*; pub const USAGE_EVENT_SCHEMA_VERSION: &str = "usage_event.v1"; @@ -31,11 +33,15 @@ pub const SYNC_ACK_V3_SCHEMA_VERSION: &str = "sync_ack.v3"; pub const SYNC_BATCH_SCHEMA_VERSION: &str = SYNC_BATCH_V3_SCHEMA_VERSION; pub const SYNC_ACK_SCHEMA_VERSION: &str = SYNC_ACK_V3_SCHEMA_VERSION; -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)] +#[derive( + Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, JsonSchema, +)] #[serde(transparent)] pub struct SourceId(pub String); -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)] +#[derive( + Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, JsonSchema, +)] #[serde(transparent)] pub struct ProviderAccountId(pub String); diff --git a/crates/statsai-core/src/quota.rs b/crates/statsai-core/src/quota.rs new file mode 100644 index 0000000..26c8238 --- /dev/null +++ b/crates/statsai-core/src/quota.rs @@ -0,0 +1,196 @@ +//! Provider-neutral quota observation and reconstructed-window contracts. + +use crate::{EventId, ProviderAccountId, SourceId, UsageCounts}; +use chrono::{DateTime, Utc}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +pub const QUOTA_OBSERVATION_SCHEMA_VERSION: &str = "quota_observation.v1"; +pub const QUOTA_WINDOW_OBSERVATION_SCHEMA_VERSION: &str = "quota_window_observation.v1"; +pub const QUOTA_WINDOW_SCHEMA_VERSION: &str = "quota_window.v1"; +pub const QUOTA_WINDOW_SYNC_PROJECTION_SCHEMA_VERSION: &str = "quota_window_sync_projection.v1"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum QuotaUsageLinkKind { + RecordEvent, + TurnEvent, + None, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum QuotaTransitionKind { + Initial, + Early, + OnOrAfterPreviousSchedule, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Default)] +pub struct QuotaCreditsV1 { + pub has_credits: Option, + pub unlimited: Option, + pub balance: Option, + pub balance_raw: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Default)] +pub struct QuotaStatusV1 { + pub plan_type: Option, + pub individual_limit: Option, + pub spend_control_state: Option, + pub reached_type: Option, + pub credits: QuotaCreditsV1, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Default)] +pub struct QuotaProjectionCreditsV1 { + pub has_credits: Option, + pub unlimited: Option, + pub balance: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Default)] +pub struct QuotaProjectionStatusV1 { + pub plan_type: Option, + pub individual_limit: Option, + pub spend_control_state: Option, + pub reached_type: Option, + pub credits: QuotaProjectionCreditsV1, +} + +impl From<&QuotaStatusV1> for QuotaProjectionStatusV1 { + fn from(status: &QuotaStatusV1) -> Self { + Self { + plan_type: status.plan_type.clone(), + individual_limit: status.individual_limit.clone(), + spend_control_state: status.spend_control_state.clone(), + reached_type: status.reached_type.clone(), + credits: QuotaProjectionCreditsV1 { + has_credits: status.credits.has_credits, + unlimited: status.credits.unlimited, + balance: status.credits.balance.clone(), + }, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct QuotaObservationV1 { + pub schema_version: String, + pub observation_id: String, + pub semantic_fingerprint: String, + pub provider: String, + pub source_id: SourceId, + pub provider_account_id: Option, + pub observed_at: DateTime, + pub source_file_path_hash: String, + pub source_record_id: String, + pub source_line_number: u64, + pub payload_hash: String, + pub usage_sample: Option, + pub usage_event_id: Option, + pub usage_link_kind: QuotaUsageLinkKind, + pub status: QuotaStatusV1, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct QuotaWindowObservationV1 { + pub schema_version: String, + pub window_observation_id: String, + pub observation_id: String, + /// Provider field name such as `primary` or `secondary`; evidence only. + pub provider_slot: String, + pub limit_id: Option, + pub window_minutes: u64, + pub used_percent: f64, + pub resets_at: DateTime, + pub resets_at_epoch_seconds: i64, +} + +/// One parsed source record plus the content-addressed provider payload. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct QuotaObservationRecordV1 { + pub observation: QuotaObservationV1, + pub windows: Vec, + pub raw_rate_limits: Value, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct QuotaChangePointV1 { + pub observed_at: DateTime, + pub used_percent: f64, + pub resets_at: DateTime, + pub resets_at_epoch_seconds: i64, + pub point_fingerprint: String, + pub provider_slot: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, Default)] +pub struct QuotaUsageTotalsV1 { + pub event_count: u64, + pub total_tokens: u64, + pub estimated_cost_micro_usd: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct QuotaWindowV1 { + pub schema_version: String, + pub window_id: String, + pub provider: String, + pub provider_account_id: Option, + /// Local source partition for unattributed evidence; omitted after account attribution. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub source_id: Option, + pub limit_id: Option, + pub window_minutes: u64, + pub inferred_start: DateTime, + pub representative_reset: DateTime, + pub representative_reset_epoch_seconds: i64, + pub reset_min: DateTime, + pub reset_min_epoch_seconds: i64, + pub reset_max: DateTime, + pub reset_max_epoch_seconds: i64, + pub first_observed_at: DateTime, + pub last_observed_at: DateTime, + pub sample_count: u64, + pub first_used_percent: f64, + pub latest_used_percent: f64, + pub minimum_used_percent: f64, + pub maximum_used_percent: f64, + pub transition: QuotaTransitionKind, + pub has_schedule_overlap: bool, + pub change_points: Vec, + pub latest_status: QuotaStatusV1, + /// Account-scoped deduplicated usage, unavailable while the window is unattributed. + pub usage_totals: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct QuotaWindowSyncProjectionV1 { + pub schema_version: String, + /// A deterministic device contribution, not a logical account-window ID. + pub projection_id: String, + pub device_id: String, + pub provider: String, + pub provider_account_id: ProviderAccountId, + pub limit_id: Option, + pub window_minutes: u64, + pub inferred_start: DateTime, + pub representative_reset: DateTime, + pub representative_reset_epoch_seconds: i64, + pub reset_min: DateTime, + pub reset_min_epoch_seconds: i64, + pub reset_max: DateTime, + pub reset_max_epoch_seconds: i64, + pub first_observed_at: DateTime, + pub last_observed_at: DateTime, + pub sample_count: u64, + pub first_used_percent: f64, + pub latest_used_percent: f64, + pub minimum_used_percent: f64, + pub maximum_used_percent: f64, + pub change_points: Vec, + pub latest_status: QuotaProjectionStatusV1, +} diff --git a/crates/statsai-store/src/archive.rs b/crates/statsai-store/src/archive.rs index 0dfdb5d..fdf1cee 100644 --- a/crates/statsai-store/src/archive.rs +++ b/crates/statsai-store/src/archive.rs @@ -5,10 +5,10 @@ use chrono::{DateTime, Utc}; use rusqlite::{params, OptionalExtension}; use serde::Serialize; use statsai_core::{ - archive_artifact_metadata_signature, ArchiveArtifactDependency, ArchiveCompleteness, + archive_artifact_metadata_signature, hash_text, ArchiveArtifactDependency, ArchiveCompleteness, ArchiveContentKind, ArchiveContentPart, ArchiveConversation, ArchiveItem, ArchiveItemKind, - ArchiveRole, CoverageStatus, ModelInfo, ProjectInfo, SourceId, TraceEdit, UsageCounts, - ARCHIVE_CONVERSATION_SCHEMA_VERSION, + ArchiveRole, CoverageStatus, ModelInfo, ProjectInfo, QuotaObservationRecordV1, SourceId, + TraceEdit, UsageCounts, ARCHIVE_CONVERSATION_SCHEMA_VERSION, }; use std::collections::{HashMap, HashSet}; use std::path::Path; @@ -21,8 +21,9 @@ use std::path::Path; // files a `Write` created become counted additions with matchable fingerprints // instead of unclassified lines; v8 binds a trace edit to the conversation its // file is written as, so edits a resumed session recorded against its parent -// stop being attributed to that parent. -const ARCHIVE_IMPORT_REVISION: &str = "archive.v8"; +// stop being attributed to that parent; v9 backfills quota observations from +// immutable Codex archives. +const ARCHIVE_IMPORT_REVISION: &str = "archive.v9"; const UNSCOPED_MISSING_CONTENT_SCOPE: &str = "unscoped"; #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize)] @@ -188,6 +189,7 @@ fn append_row_placeholders(sql: &mut String, rows: usize, columns: usize) { } impl Store { + #[allow(clippy::too_many_arguments)] pub fn store_archive_scan_with_code_changes( &self, source_id: &SourceId, @@ -196,6 +198,7 @@ impl Store { artifact_dependencies: &[ArchiveArtifactDependency], trace_edits: &[TraceEdit], trace_coverage: CoverageStatus, + quota_observations: &[QuotaObservationRecordV1], ) -> Result { self.with_immediate_transaction(|| { let result = self.upsert_archive_conversations(conversations)?; @@ -211,6 +214,15 @@ impl Store { imported_entries, artifact_dependencies, )?; + let source_file_path_hashes = imported_entries + .iter() + .map(|entry| hash_text(&entry.cache_key)) + .collect::>(); + self.replace_quota_observations_for_source_files_inner( + source_id, + &source_file_path_hashes, + quota_observations, + )?; Ok(result) }) } @@ -297,7 +309,19 @@ impl Store { "DELETE FROM archive_import_state WHERE source_id = ?1 AND cache_key = ?2", params![&source_id.0, cache_key], )?; + self.conn.execute( + "DELETE FROM quota_window_observations WHERE observation_id IN (SELECT observation_id FROM quota_observations WHERE source_id = ?1 AND source_file_path_hash = ?2)", + params![&source_id.0, hash_text(cache_key)], + )?; + self.conn.execute( + "DELETE FROM quota_observations WHERE source_id = ?1 AND source_file_path_hash = ?2", + params![&source_id.0, hash_text(cache_key)], + )?; } + self.conn.execute( + "DELETE FROM quota_payloads WHERE payload_hash NOT IN (SELECT payload_hash FROM quota_observations)", + [], + )?; Ok(removed_cache_keys.len() as u64) }) } @@ -1602,6 +1626,7 @@ mod tests { &[], &[], CoverageStatus::Unavailable, + &[], ) .expect("store archive scan"); assert!( @@ -1662,6 +1687,7 @@ mod tests { std::slice::from_ref(&dependency), &[], CoverageStatus::Unavailable, + &[], ) .expect("store archive scan"); assert!(store diff --git a/crates/statsai-store/src/lib.rs b/crates/statsai-store/src/lib.rs index e91467b..e6ff279 100644 --- a/crates/statsai-store/src/lib.rs +++ b/crates/statsai-store/src/lib.rs @@ -4,6 +4,7 @@ mod archive; mod code_changes; mod migrations; mod privacy; +mod quota; mod snapshot; mod tasks; @@ -48,6 +49,7 @@ pub use privacy::{ FilteredConversationMetadata, FilteredConversationRecord, PrivacyDatasetStatus, PrivacyFailureRecord, PrivacyFindingRecord, }; +pub use quota::{QuotaDateRange, QuotaQuery, QuotaStatus}; pub use tasks::{ derive_task_work_items, NamedTaskBenchmark, TaskBenchmarkMetrics, TaskBenchmarkReport, TaskDeletionImpact, TaskRebuildReport, TaskRebuildTimings, TaskStats, @@ -4813,6 +4815,7 @@ fn reattribute_source_records(store: &Store, source_id: &SourceId) -> Result<()> } store.rewrite_events(&events)?; store.rewrite_summaries(&summaries)?; + store.reattribute_quota_observations(source_id)?; Ok(()) } diff --git a/crates/statsai-store/src/migrations.rs b/crates/statsai-store/src/migrations.rs index c6bc838..6e2d01d 100644 --- a/crates/statsai-store/src/migrations.rs +++ b/crates/statsai-store/src/migrations.rs @@ -2,7 +2,7 @@ use anyhow::{bail, Context, Result}; use chrono::Utc; use rusqlite::{Connection, OptionalExtension}; -pub const CURRENT_SCHEMA_VERSION: i64 = 17; +pub const CURRENT_SCHEMA_VERSION: i64 = 18; pub fn migrate(conn: &Connection) -> Result<()> { if let Some(current) = existing_schema_version(conn)? { @@ -136,6 +136,7 @@ fn apply_migration(conn: &Connection, version: i64) -> Result<()> { 15 => apply_migration_015(conn), 16 => apply_migration_016(conn), 17 => apply_migration_017(conn), + 18 => apply_migration_018(conn), _ => bail!("unsupported schema migration version {version}"), } } @@ -702,6 +703,57 @@ fn apply_migration_017(conn: &Connection) -> Result<()> { Ok(()) } +fn apply_migration_018(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" + CREATE TABLE IF NOT EXISTS quota_payloads ( + payload_hash TEXT PRIMARY KEY, + provider TEXT NOT NULL, + payload TEXT NOT NULL, + created_at TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS quota_observations ( + observation_id TEXT PRIMARY KEY, + semantic_fingerprint TEXT NOT NULL, + provider TEXT NOT NULL, + source_id TEXT NOT NULL, + provider_account_id TEXT, + observed_at TEXT NOT NULL, + source_file_path_hash TEXT NOT NULL, + source_record_id TEXT NOT NULL, + usage_event_id TEXT, + usage_link_kind TEXT NOT NULL, + payload_hash TEXT NOT NULL, + payload TEXT NOT NULL, + FOREIGN KEY (payload_hash) REFERENCES quota_payloads(payload_hash) + ); + CREATE INDEX IF NOT EXISTS quota_observations_scope_idx + ON quota_observations (provider, provider_account_id, observed_at, observation_id); + CREATE INDEX IF NOT EXISTS quota_observations_source_idx + ON quota_observations (source_id, source_file_path_hash, observed_at, observation_id); + CREATE INDEX IF NOT EXISTS quota_observations_semantic_idx + ON quota_observations (semantic_fingerprint, observation_id); + + CREATE TABLE IF NOT EXISTS quota_window_observations ( + window_observation_id TEXT PRIMARY KEY, + observation_id TEXT NOT NULL, + provider_slot TEXT NOT NULL, + limit_id TEXT, + window_minutes INTEGER NOT NULL, + used_percent REAL NOT NULL, + resets_at INTEGER NOT NULL, + payload TEXT NOT NULL, + FOREIGN KEY (observation_id) REFERENCES quota_observations(observation_id) + ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS quota_window_observations_reconstruct_idx + ON quota_window_observations (window_minutes, limit_id, resets_at, observation_id); + "#, + )?; + Ok(()) +} + fn ensure_local_task_tables(conn: &Connection) -> Result<()> { conn.execute_batch( r#" @@ -874,20 +926,20 @@ mod tests { applied_at TEXT NOT NULL ); INSERT INTO schema_migrations (version, applied_at) - VALUES (18, '2026-08-19T00:00:00Z'); + VALUES (19, '2026-08-19T00:00:00Z'); "#, ) .expect("create future schema marker"); - let error = migrate(&conn).expect_err("schema 18 must be rejected by schema 17 binary"); + let error = migrate(&conn).expect_err("schema 19 must be rejected by schema 18 binary"); assert_eq!( error.to_string(), - "database schema version 18 is newer than this StatsAI binary supports (17); upgrade StatsAI or use a compatible database" + "database schema version 19 is newer than this StatsAI binary supports (18); upgrade StatsAI or use a compatible database" ); assert_eq!( current_schema_version(&conn).expect("read unchanged version"), - 18 + 19 ); } diff --git a/crates/statsai-store/src/quota.rs b/crates/statsai-store/src/quota.rs new file mode 100644 index 0000000..b88a5d0 --- /dev/null +++ b/crates/statsai-store/src/quota.rs @@ -0,0 +1,1905 @@ +use super::{assignment_for_timestamp, Store}; +use anyhow::Result; +use chrono::{DateTime, Duration, Utc}; +use rusqlite::{params, OptionalExtension}; +use serde::Serialize; +use statsai_core::{ + hash_text, periods_overlap, EventId, ProviderAccountId, QuotaChangePointV1, + QuotaObservationRecordV1, QuotaObservationV1, QuotaProjectionStatusV1, QuotaTransitionKind, + QuotaUsageLinkKind, QuotaUsageTotalsV1, QuotaWindowObservationV1, QuotaWindowSyncProjectionV1, + QuotaWindowV1, SourceAccountAssignment, SourceId, QUOTA_WINDOW_SCHEMA_VERSION, + QUOTA_WINDOW_SYNC_PROJECTION_SCHEMA_VERSION, +}; +use std::collections::{BTreeMap, HashMap, HashSet}; + +const RESET_CLUSTER_TOLERANCE_SECONDS: i64 = 5 * 60; +pub const SYNC_ELIGIBLE_WINDOW_MINUTES: u64 = 10_080; + +#[derive(Debug, Clone, Default)] +pub struct QuotaQuery { + pub provider: Option, + pub provider_account_id: Option, + pub source_id: Option, + pub from: Option>, + pub to: Option>, + pub limit_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct QuotaDateRange { + pub first: DateTime, + pub last: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct QuotaStatus { + pub schema_version: String, + pub total_observations: u64, + pub distinct_observations: u64, + pub duplicate_observations: u64, + pub attributed_observations: u64, + pub unattributed_observations: u64, + pub attributed_range: Option, + pub unattributed_range: Option, + pub weekly_observations: u64, + pub weekly_sync_eligible_observations: u64, + pub weekly_sync_eligible_coverage_percent: f64, + pub assignment_overlap_warnings: Vec, +} + +impl Store { + pub fn clear_orphaned_quota_usage_links(&self) -> Result { + Ok(self.conn.execute( + r#" + UPDATE quota_observations + SET usage_event_id = NULL, + usage_link_kind = 'none', + payload = json_set( + payload, + '$.usage_event_id', NULL, + '$.usage_link_kind', 'none' + ) + WHERE usage_event_id IS NOT NULL + AND NOT EXISTS ( + SELECT 1 FROM usage_events e + WHERE e.event_id = quota_observations.usage_event_id + ) + "#, + [], + )? as u64) + } + + pub fn upsert_quota_observations(&self, records: &[QuotaObservationRecordV1]) -> Result { + self.with_immediate_transaction(|| self.upsert_quota_observations_inner(records)) + } + + pub(crate) fn upsert_quota_observations_inner( + &self, + records: &[QuotaObservationRecordV1], + ) -> Result { + let mut written = 0u64; + let mut assignments_by_source = HashMap::new(); + for record in records { + let source_id = &record.observation.source_id; + if !assignments_by_source.contains_key(source_id) { + assignments_by_source.insert( + source_id.clone(), + self.list_source_account_assignments_for_source(source_id)?, + ); + } + let mut record = record.clone(); + let assignments = assignments_by_source + .get(source_id) + .expect("source assignments inserted above"); + record.observation.provider_account_id = + assignment_for_timestamp(assignments, record.observation.observed_at) + .map(|assignment| assignment.provider_account_id.clone()); + + if record.observation.usage_event_id.is_none() { + let existing = self + .conn + .query_row( + "SELECT usage_event_id, usage_link_kind, payload FROM quota_observations WHERE observation_id = ?1", + [&record.observation.observation_id], + |row| { + Ok(( + row.get::<_, Option>(0)?, + row.get::<_, String>(1)?, + row.get::<_, String>(2)?, + )) + }, + ) + .optional()?; + if let Some((Some(event_id), link_kind, payload)) = existing { + let existing_observation: QuotaObservationV1 = serde_json::from_str(&payload)?; + if matching_positive_usage_sample(&record.observation, &existing_observation) { + record.observation.usage_event_id = Some(EventId(event_id)); + record.observation.usage_link_kind = parse_usage_link_kind(&link_kind); + } + } + } + + self.conn.execute( + r#" + INSERT OR IGNORE INTO quota_payloads + (payload_hash, provider, payload, created_at) + VALUES (?1, ?2, ?3, ?4) + "#, + params![ + &record.observation.payload_hash, + &record.observation.provider, + serde_json::to_string(&record.raw_rate_limits)?, + Utc::now().to_rfc3339(), + ], + )?; + let observation_payload = serde_json::to_string(&record.observation)?; + self.conn.execute( + r#" + INSERT INTO quota_observations ( + observation_id, semantic_fingerprint, provider, source_id, + provider_account_id, observed_at, source_file_path_hash, + source_record_id, usage_event_id, usage_link_kind, payload_hash, payload + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12) + ON CONFLICT(observation_id) DO UPDATE SET + semantic_fingerprint = excluded.semantic_fingerprint, + provider = excluded.provider, + source_id = excluded.source_id, + provider_account_id = excluded.provider_account_id, + observed_at = excluded.observed_at, + source_file_path_hash = excluded.source_file_path_hash, + source_record_id = excluded.source_record_id, + usage_event_id = excluded.usage_event_id, + usage_link_kind = excluded.usage_link_kind, + payload_hash = excluded.payload_hash, + payload = excluded.payload + "#, + params![ + &record.observation.observation_id, + &record.observation.semantic_fingerprint, + &record.observation.provider, + &record.observation.source_id.0, + record + .observation + .provider_account_id + .as_ref() + .map(|id| id.0.as_str()), + record.observation.observed_at.to_rfc3339(), + &record.observation.source_file_path_hash, + &record.observation.source_record_id, + record + .observation + .usage_event_id + .as_ref() + .map(|id| id.0.as_str()), + usage_link_kind_label(record.observation.usage_link_kind), + &record.observation.payload_hash, + observation_payload, + ], + )?; + self.conn.execute( + "DELETE FROM quota_window_observations WHERE observation_id = ?1", + [&record.observation.observation_id], + )?; + for window in &record.windows { + self.conn.execute( + r#" + INSERT INTO quota_window_observations ( + window_observation_id, observation_id, provider_slot, limit_id, + window_minutes, used_percent, resets_at, payload + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + "#, + params![ + &window.window_observation_id, + &window.observation_id, + &window.provider_slot, + window.limit_id.as_deref(), + i64::try_from(window.window_minutes).unwrap_or(i64::MAX), + window.used_percent, + window.resets_at_epoch_seconds, + serde_json::to_string(window)?, + ], + )?; + } + written = written.saturating_add(1); + } + Ok(written) + } + + pub(crate) fn replace_quota_observations_for_source_files_inner( + &self, + source_id: &SourceId, + source_file_path_hashes: &[String], + records: &[QuotaObservationRecordV1], + ) -> Result { + let retained_observation_ids = records + .iter() + .filter(|record| { + record.observation.source_id == *source_id + && source_file_path_hashes.contains(&record.observation.source_file_path_hash) + }) + .map(|record| record.observation.observation_id.as_str()) + .collect::>(); + + for source_file_path_hash in source_file_path_hashes { + let mut statement = self.conn.prepare( + "SELECT observation_id FROM quota_observations + WHERE source_id = ?1 AND source_file_path_hash = ?2", + )?; + let stored_observation_ids = statement + .query_map(params![&source_id.0, source_file_path_hash], |row| { + row.get::<_, String>(0) + })? + .collect::>>()?; + drop(statement); + + for observation_id in stored_observation_ids { + if retained_observation_ids.contains(observation_id.as_str()) { + continue; + } + self.conn.execute( + "DELETE FROM quota_window_observations WHERE observation_id = ?1", + [&observation_id], + )?; + self.conn.execute( + "DELETE FROM quota_observations WHERE observation_id = ?1", + [&observation_id], + )?; + } + } + + let written = self.upsert_quota_observations_inner(records)?; + self.delete_unreferenced_quota_payloads()?; + Ok(written) + } + + pub fn quota_observations( + &self, + query: &QuotaQuery, + collapse_duplicates: bool, + ) -> Result> { + let mut observations = BTreeMap::::new(); + let mut statement = self.conn.prepare( + r#" + SELECT q.payload, q.provider_account_id, q.usage_event_id, q.usage_link_kind, + p.payload + FROM quota_observations q + JOIN quota_payloads p ON p.payload_hash = q.payload_hash + ORDER BY q.observed_at, q.observation_id + "#, + )?; + let rows = statement.query_map([], |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, Option>(1)?, + row.get::<_, Option>(2)?, + row.get::<_, String>(3)?, + row.get::<_, String>(4)?, + )) + })?; + for row in rows { + let (payload, account_id, usage_event_id, usage_link_kind, raw_payload) = row?; + let mut observation: QuotaObservationV1 = serde_json::from_str(&payload)?; + observation.provider_account_id = account_id.map(ProviderAccountId); + observation.usage_event_id = usage_event_id.map(EventId); + observation.usage_link_kind = parse_usage_link_kind(&usage_link_kind); + if !observation_matches_query(&observation, query) { + continue; + } + observations.insert( + observation.observation_id.clone(), + QuotaObservationRecordV1 { + observation, + windows: Vec::new(), + raw_rate_limits: serde_json::from_str(&raw_payload)?, + }, + ); + } + drop(statement); + + let mut statement = self.conn.prepare( + "SELECT observation_id, payload FROM quota_window_observations ORDER BY resets_at, window_observation_id", + )?; + let rows = statement.query_map([], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) + })?; + for row in rows { + let (observation_id, payload) = row?; + let Some(record) = observations.get_mut(&observation_id) else { + continue; + }; + let window: QuotaWindowObservationV1 = serde_json::from_str(&payload)?; + if query + .limit_id + .as_deref() + .is_some_and(|limit| window.limit_id.as_deref() != Some(limit)) + { + continue; + } + record.windows.push(window); + } + let mut records = observations.into_values().collect::>(); + records.sort_by_key(|record| { + ( + record.observation.observed_at, + record.observation.observation_id.clone(), + ) + }); + Ok(if collapse_duplicates { + collapse_semantic_duplicates(records) + } else { + records + }) + } + + pub fn delete_quota_observations_for_sources(&self, source_ids: &[SourceId]) -> Result { + self.with_immediate_transaction(|| { + let mut deleted = 0u64; + for source_id in source_ids { + self.conn.execute( + "DELETE FROM quota_window_observations WHERE observation_id IN (SELECT observation_id FROM quota_observations WHERE source_id = ?1)", + [&source_id.0], + )?; + deleted = deleted.saturating_add(self.conn.execute( + "DELETE FROM quota_observations WHERE source_id = ?1", + [&source_id.0], + )? as u64); + } + self.delete_unreferenced_quota_payloads()?; + Ok(deleted) + }) + } + + pub fn delete_quota_observations_for_source_file_hashes( + &self, + source_id: &SourceId, + file_hashes: &[String], + ) -> Result { + self.with_immediate_transaction(|| { + let mut deleted = 0u64; + for file_hash in file_hashes { + self.conn.execute( + "DELETE FROM quota_window_observations WHERE observation_id IN (SELECT observation_id FROM quota_observations WHERE source_id = ?1 AND source_file_path_hash = ?2)", + params![&source_id.0, file_hash], + )?; + deleted = deleted.saturating_add(self.conn.execute( + "DELETE FROM quota_observations WHERE source_id = ?1 AND source_file_path_hash = ?2", + params![&source_id.0, file_hash], + )? as u64); + } + self.delete_unreferenced_quota_payloads()?; + Ok(deleted) + }) + } + + fn delete_unreferenced_quota_payloads(&self) -> Result<()> { + self.conn.execute( + "DELETE FROM quota_payloads WHERE payload_hash NOT IN (SELECT payload_hash FROM quota_observations)", + [], + )?; + Ok(()) + } + + pub fn reattribute_quota_observations(&self, source_id: &SourceId) -> Result { + let assignments = self.list_source_account_assignments_for_source(source_id)?; + let records = self.quota_observations_for_source(source_id)?; + self.with_immediate_transaction(|| { + let mut changed = 0u64; + for mut record in records { + let account_id = assignment_for_timestamp( + &assignments, + record.observation.observed_at, + ) + .map(|assignment| assignment.provider_account_id.clone()); + if record.observation.provider_account_id == account_id { + continue; + } + record.observation.provider_account_id = account_id; + changed = changed.saturating_add(self.conn.execute( + "UPDATE quota_observations SET provider_account_id = ?2, payload = ?3 WHERE observation_id = ?1", + params![ + &record.observation.observation_id, + record.observation.provider_account_id.as_ref().map(|id| id.0.as_str()), + serde_json::to_string(&record.observation)?, + ], + )? as u64); + } + Ok(changed) + }) + } + + fn quota_observations_for_source( + &self, + source_id: &SourceId, + ) -> Result> { + self.quota_observations( + &QuotaQuery { + source_id: Some(source_id.clone()), + ..QuotaQuery::default() + }, + false, + ) + } + + pub fn quota_status(&self, query: &QuotaQuery) -> Result { + let raw = self.quota_observations(query, false)?; + let total_observations = raw.len(); + let distinct = collapse_semantic_duplicates(raw); + let attributed = distinct + .iter() + .filter(|record| record.observation.provider_account_id.is_some()) + .collect::>(); + let unattributed = distinct + .iter() + .filter(|record| record.observation.provider_account_id.is_none()) + .collect::>(); + let weekly = distinct + .iter() + .filter(|record| { + record + .windows + .iter() + .any(|window| window.window_minutes >= SYNC_ELIGIBLE_WINDOW_MINUTES) + }) + .collect::>(); + let weekly_eligible = weekly + .iter() + .filter(|record| record.observation.provider_account_id.is_some()) + .count() as u64; + let assignments = self + .list_source_account_assignments()? + .into_iter() + .filter(|assignment| assignment_matches_query(assignment, query)) + .collect::>(); + let mut warnings = Vec::new(); + for (index, left) in assignments.iter().enumerate() { + for right in assignments.iter().skip(index + 1) { + if left.source_id == right.source_id + && periods_overlap( + left.started_at, + left.ended_at, + right.started_at, + right.ended_at, + ) + { + warnings.push(format!( + "source {} assignments {} and {} overlap", + left.source_id.0, left.assignment_id.0, right.assignment_id.0 + )); + } + } + } + Ok(QuotaStatus { + schema_version: "quota_status.v1".to_string(), + total_observations: total_observations as u64, + distinct_observations: distinct.len() as u64, + duplicate_observations: total_observations.saturating_sub(distinct.len()) as u64, + attributed_observations: attributed.len() as u64, + unattributed_observations: unattributed.len() as u64, + attributed_range: observation_range(&attributed), + unattributed_range: observation_range(&unattributed), + weekly_observations: weekly.len() as u64, + weekly_sync_eligible_observations: weekly_eligible, + weekly_sync_eligible_coverage_percent: if weekly.is_empty() { + 100.0 + } else { + weekly_eligible as f64 * 100.0 / weekly.len() as f64 + }, + assignment_overlap_warnings: warnings, + }) + } + + pub fn quota_windows(&self, query: &QuotaQuery) -> Result> { + let mut windows = self.quota_windows_without_usage_totals(query)?; + self.enrich_quota_window_usage_totals(&mut windows)?; + Ok(windows) + } + + pub fn quota_windows_without_usage_totals( + &self, + query: &QuotaQuery, + ) -> Result> { + let mut reconstruction_query = query.clone(); + reconstruction_query.from = None; + reconstruction_query.to = None; + let records = self.quota_observations(&reconstruction_query, true)?; + let mut grouped = BTreeMap::>::new(); + for record in records { + for window in &record.windows { + grouped + .entry(WindowScope { + provider: record.observation.provider.clone(), + provider_account_id: record.observation.provider_account_id.clone(), + source_id: record + .observation + .provider_account_id + .is_none() + .then(|| record.observation.source_id.clone()), + limit_id: window.limit_id.clone(), + window_minutes: window.window_minutes, + }) + .or_default() + .push(WindowPoint { + observation: record.observation.clone(), + window: window.clone(), + }); + } + } + + let mut windows = Vec::new(); + for (scope, mut points) in grouped { + points.sort_by_key(|point| { + ( + point.window.resets_at_epoch_seconds, + point.observation.observed_at, + ) + }); + let mut clusters: Vec> = Vec::new(); + for point in points { + if let Some(cluster) = clusters.last_mut() { + let min_reset = cluster + .first() + .map(|point| point.window.resets_at_epoch_seconds) + .unwrap_or(point.window.resets_at_epoch_seconds); + if point.window.resets_at_epoch_seconds - min_reset + <= RESET_CLUSTER_TOLERANCE_SECONDS + { + cluster.push(point); + continue; + } + } + clusters.push(vec![point]); + } + + let schedules = clusters + .iter() + .map(|cluster| QuotaClusterSchedule::from_points(&scope, cluster)) + .collect::>(); + let mut scope_windows = Vec::with_capacity(clusters.len()); + for (index, (cluster, schedule)) in + clusters.into_iter().zip(schedules.iter()).enumerate() + { + if !quota_cluster_matches_time_query(schedule, query) { + continue; + } + let mut window = Self::materialize_quota_window(&scope, schedule, cluster); + window.transition = if index == 0 { + QuotaTransitionKind::Initial + } else if schedule.first_observed_at < schedules[index - 1].representative_reset { + QuotaTransitionKind::Early + } else { + QuotaTransitionKind::OnOrAfterPreviousSchedule + }; + window.has_schedule_overlap = schedules.iter().enumerate().any(|(other, peer)| { + other != index + && schedule.inferred_start < peer.representative_reset + && schedule.representative_reset > peer.inferred_start + }); + scope_windows.push(window); + } + windows.extend(scope_windows); + } + windows.sort_by(|left, right| { + right + .representative_reset + .cmp(&left.representative_reset) + .then_with(|| right.window_minutes.cmp(&left.window_minutes)) + .then_with(|| left.window_id.cmp(&right.window_id)) + }); + Ok(windows) + } + + fn materialize_quota_window( + scope: &WindowScope, + schedule: &QuotaClusterSchedule, + mut points: Vec, + ) -> QuotaWindowV1 { + points.sort_by_key(|point| point.observation.observed_at); + let mut change_points = Vec::new(); + for point in &points { + let changed = change_points + .last() + .is_none_or(|previous: &QuotaChangePointV1| { + previous.used_percent != point.window.used_percent + || previous.resets_at_epoch_seconds != point.window.resets_at_epoch_seconds + }); + if !changed { + continue; + } + let point_fingerprint = quota_point_fingerprint(scope, point); + change_points.push(QuotaChangePointV1 { + observed_at: point.observation.observed_at, + used_percent: point.window.used_percent, + resets_at: point.window.resets_at, + resets_at_epoch_seconds: point.window.resets_at_epoch_seconds, + point_fingerprint, + provider_slot: point.window.provider_slot.clone(), + }); + } + let first = points.first().expect("non-empty reset cluster"); + let latest = points.last().expect("non-empty reset cluster"); + let anchor = change_points + .first() + .expect("non-empty reset cluster has a change point"); + let identity_scope = scope.provider_account_id.as_ref().map_or_else( + || { + format!( + "unassigned:{}", + scope + .source_id + .as_ref() + .map(|id| id.0.as_str()) + .unwrap_or("unknown") + ) + }, + |id| id.0.clone(), + ); + let window_id = format!( + "quota_window_{}", + &hash_text(&format!( + "quota_window.v1:{}:{identity_scope}:{}:{}:{}", + scope.provider, + scope.limit_id.as_deref().unwrap_or("default"), + scope.window_minutes, + anchor.point_fingerprint + ))[..32] + ); + QuotaWindowV1 { + schema_version: QUOTA_WINDOW_SCHEMA_VERSION.to_string(), + window_id, + provider: scope.provider.clone(), + provider_account_id: scope.provider_account_id.clone(), + source_id: scope.source_id.clone(), + limit_id: scope.limit_id.clone(), + window_minutes: scope.window_minutes, + inferred_start: schedule.inferred_start, + representative_reset: schedule.representative_reset, + representative_reset_epoch_seconds: schedule.representative_reset_epoch_seconds, + reset_min: schedule.reset_min, + reset_min_epoch_seconds: schedule.reset_min_epoch_seconds, + reset_max: schedule.reset_max, + reset_max_epoch_seconds: schedule.reset_max_epoch_seconds, + first_observed_at: first.observation.observed_at, + last_observed_at: latest.observation.observed_at, + sample_count: points.len() as u64, + first_used_percent: first.window.used_percent, + latest_used_percent: latest.window.used_percent, + minimum_used_percent: points + .iter() + .map(|point| point.window.used_percent) + .fold(f64::INFINITY, f64::min), + maximum_used_percent: points + .iter() + .map(|point| point.window.used_percent) + .fold(f64::NEG_INFINITY, f64::max), + transition: QuotaTransitionKind::Initial, + has_schedule_overlap: false, + change_points, + latest_status: latest.observation.status.clone(), + usage_totals: None, + } + } + + pub fn enrich_quota_window_usage_totals(&self, windows: &mut [QuotaWindowV1]) -> Result<()> { + let mut windows_by_account = + HashMap::>>::new(); + for (index, window) in windows.iter_mut().enumerate() { + let Some(account_id) = window.provider_account_id.clone() else { + window.usage_totals = None; + continue; + }; + window.usage_totals = Some(QuotaUsageTotalsV1::default()); + windows_by_account + .entry(window.provider.clone()) + .or_default() + .entry(account_id) + .or_default() + .push(index); + } + if windows_by_account.is_empty() { + return Ok(()); + } + + let first_start = windows + .iter() + .filter(|window| window.provider_account_id.is_some()) + .map(|window| window.inferred_start) + .min() + .expect("attributed window index is non-empty"); + let last_reset = windows + .iter() + .filter(|window| window.provider_account_id.is_some()) + .map(|window| window.representative_reset) + .max() + .expect("attributed window index is non-empty"); + + for event in self.events_in_period(Some(first_start), last_reset)? { + let Some(account_id) = event.provider_account_id.as_ref() else { + continue; + }; + let Some(indexes) = windows_by_account + .get(event.provider.as_str()) + .and_then(|accounts| accounts.get(account_id)) + else { + continue; + }; + for &index in indexes { + let window = &mut windows[index]; + if event.session.started_at < window.inferred_start + || event.session.started_at >= window.representative_reset + { + continue; + } + let totals = window + .usage_totals + .as_mut() + .expect("attributed windows receive totals above"); + totals.event_count = totals.event_count.saturating_add(1); + totals.total_tokens = totals + .total_tokens + .saturating_add(event.usage.computed_total()); + if let Some(value) = event + .cost + .provider_reported_micro_usd_value() + .or_else(|| event.cost.estimated_micro_usd()) + { + totals.estimated_cost_micro_usd = Some( + totals + .estimated_cost_micro_usd + .unwrap_or_default() + .saturating_add(value), + ); + } + } + } + Ok(()) + } + + pub fn quota_sync_projections( + &self, + query: &QuotaQuery, + device_id: &str, + ) -> Result> { + Ok(self + .quota_windows_without_usage_totals(query)? + .into_iter() + .filter(|window| window.window_minutes >= SYNC_ELIGIBLE_WINDOW_MINUTES) + .filter_map(|window| { + let provider_account_id = window.provider_account_id.clone()?; + let anchor = window.change_points.first()?; + let projection_id = format!( + "quota_projection_{}", + &hash_text(&format!( + "quota_projection.v1:{device_id}:{}:{}:{}:{}:{}", + window.provider, + provider_account_id.0, + window.limit_id.as_deref().unwrap_or("default"), + window.window_minutes, + anchor.point_fingerprint + ))[..32] + ); + Some(QuotaWindowSyncProjectionV1 { + schema_version: QUOTA_WINDOW_SYNC_PROJECTION_SCHEMA_VERSION.to_string(), + projection_id, + device_id: device_id.to_string(), + provider: window.provider, + provider_account_id, + limit_id: window.limit_id, + window_minutes: window.window_minutes, + inferred_start: window.inferred_start, + representative_reset: window.representative_reset, + representative_reset_epoch_seconds: window.representative_reset_epoch_seconds, + reset_min: window.reset_min, + reset_min_epoch_seconds: window.reset_min_epoch_seconds, + reset_max: window.reset_max, + reset_max_epoch_seconds: window.reset_max_epoch_seconds, + first_observed_at: window.first_observed_at, + last_observed_at: window.last_observed_at, + sample_count: window.sample_count, + first_used_percent: window.first_used_percent, + latest_used_percent: window.latest_used_percent, + minimum_used_percent: window.minimum_used_percent, + maximum_used_percent: window.maximum_used_percent, + change_points: window.change_points, + latest_status: QuotaProjectionStatusV1::from(&window.latest_status), + }) + }) + .collect()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +struct WindowScope { + provider: String, + provider_account_id: Option, + source_id: Option, + limit_id: Option, + window_minutes: u64, +} + +#[derive(Debug, Clone)] +struct WindowPoint { + observation: QuotaObservationV1, + window: QuotaWindowObservationV1, +} + +#[derive(Debug, Clone)] +struct QuotaClusterSchedule { + inferred_start: DateTime, + representative_reset: DateTime, + representative_reset_epoch_seconds: i64, + reset_min: DateTime, + reset_min_epoch_seconds: i64, + reset_max: DateTime, + reset_max_epoch_seconds: i64, + first_observed_at: DateTime, + last_observed_at: DateTime, +} + +impl QuotaClusterSchedule { + fn from_points(scope: &WindowScope, points: &[WindowPoint]) -> Self { + let first = points.first().expect("non-empty reset cluster"); + let last = points.last().expect("non-empty reset cluster"); + let representative = &points[points.len() / 2]; + let representative_reset = representative.window.resets_at; + Self { + inferred_start: representative_reset + - Duration::minutes(i64::try_from(scope.window_minutes).unwrap_or(i64::MAX)), + representative_reset, + representative_reset_epoch_seconds: representative.window.resets_at_epoch_seconds, + reset_min: first.window.resets_at, + reset_min_epoch_seconds: first.window.resets_at_epoch_seconds, + reset_max: last.window.resets_at, + reset_max_epoch_seconds: last.window.resets_at_epoch_seconds, + first_observed_at: points + .iter() + .map(|point| point.observation.observed_at) + .min() + .expect("non-empty reset cluster"), + last_observed_at: points + .iter() + .map(|point| point.observation.observed_at) + .max() + .expect("non-empty reset cluster"), + } + } +} + +fn quota_point_fingerprint(scope: &WindowScope, point: &WindowPoint) -> String { + let identity_scope = scope.provider_account_id.as_ref().map_or_else( + || { + format!( + "unassigned:{}", + scope + .source_id + .as_ref() + .map(|id| id.0.as_str()) + .unwrap_or("unknown") + ) + }, + |id| id.0.clone(), + ); + hash_text(&format!( + "quota_point.v1:{}:{}:{}:{}:{}:{}:{}", + scope.provider, + identity_scope, + scope.limit_id.as_deref().unwrap_or("default"), + scope.window_minutes, + point.observation.observed_at.to_rfc3339(), + point.window.used_percent, + point.window.resets_at_epoch_seconds, + )) +} + +fn observation_matches_query(observation: &QuotaObservationV1, query: &QuotaQuery) -> bool { + query + .provider + .as_deref() + .is_none_or(|provider| observation.provider == provider) + && query + .provider_account_id + .as_ref() + .is_none_or(|account| observation.provider_account_id.as_ref() == Some(account)) + && query + .source_id + .as_ref() + .is_none_or(|source_id| &observation.source_id == source_id) + && query + .from + .is_none_or(|from| observation.observed_at >= from) + && query.to.is_none_or(|to| observation.observed_at <= to) +} + +fn assignment_matches_query(assignment: &SourceAccountAssignment, query: &QuotaQuery) -> bool { + query + .provider + .as_deref() + .is_none_or(|provider| assignment.provider == provider) + && query + .provider_account_id + .as_ref() + .is_none_or(|account| &assignment.provider_account_id == account) + && query + .source_id + .as_ref() + .is_none_or(|source_id| &assignment.source_id == source_id) + && query + .from + .is_none_or(|from| assignment.ended_at.is_none_or(|ended_at| ended_at > from)) + && query.to.is_none_or(|to| assignment.started_at <= to) +} + +fn quota_cluster_matches_time_query(schedule: &QuotaClusterSchedule, query: &QuotaQuery) -> bool { + query + .from + .is_none_or(|from| schedule.last_observed_at >= from) + && query.to.is_none_or(|to| schedule.first_observed_at <= to) +} + +fn collapse_semantic_duplicates( + records: Vec, +) -> Vec { + let mut semantic_groups = HashMap::>::new(); + for record in records { + semantic_groups + .entry(record.observation.semantic_fingerprint.clone()) + .or_default() + .push(record); + } + + let mut collapsed = Vec::new(); + for group in semantic_groups.into_values() { + let attributed_accounts = group + .iter() + .filter_map(|record| record.observation.provider_account_id.as_ref()) + .collect::>(); + if attributed_accounts.is_empty() { + let mut source_scoped = HashMap::::new(); + for record in group { + let source_id = record.observation.source_id.clone(); + match source_scoped.get_mut(&source_id) { + Some(existing) + if observation_quality(&record) > observation_quality(existing) => + { + *existing = record; + } + Some(_) => {} + None => { + source_scoped.insert(source_id, record); + } + } + } + collapsed.extend(source_scoped.into_values()); + continue; + } + if attributed_accounts.len() == 1 { + if let Some(best) = group.into_iter().max_by_key(observation_quality) { + collapsed.push(best); + } + continue; + } + + let mut account_scoped = + HashMap::, QuotaObservationRecordV1>::new(); + for record in group { + let account_id = record.observation.provider_account_id.clone(); + if account_id.is_none() { + continue; + } + match account_scoped.get_mut(&account_id) { + Some(existing) if observation_quality(&record) > observation_quality(existing) => { + *existing = record; + } + Some(_) => {} + None => { + account_scoped.insert(account_id, record); + } + } + } + collapsed.extend(account_scoped.into_values()); + } + + collapsed.sort_by_key(|record| { + ( + record.observation.observed_at, + record.observation.observation_id.clone(), + ) + }); + collapsed +} + +fn observation_quality(record: &QuotaObservationRecordV1) -> (bool, bool, usize) { + ( + record.observation.provider_account_id.is_some(), + record.observation.usage_event_id.is_some(), + record.windows.len(), + ) +} + +fn matching_positive_usage_sample( + incoming: &QuotaObservationV1, + existing: &QuotaObservationV1, +) -> bool { + incoming.usage_sample.as_ref().is_some_and(|sample| { + sample.computed_total() > 0 && existing.usage_sample.as_ref() == Some(sample) + }) +} + +fn observation_range(records: &[&QuotaObservationRecordV1]) -> Option { + Some(QuotaDateRange { + first: records + .iter() + .map(|record| record.observation.observed_at) + .min()?, + last: records + .iter() + .map(|record| record.observation.observed_at) + .max()?, + }) +} + +fn usage_link_kind_label(kind: QuotaUsageLinkKind) -> &'static str { + match kind { + QuotaUsageLinkKind::RecordEvent => "record_event", + QuotaUsageLinkKind::TurnEvent => "turn_event", + QuotaUsageLinkKind::None => "none", + } +} + +fn parse_usage_link_kind(value: &str) -> QuotaUsageLinkKind { + match value { + "record_event" => QuotaUsageLinkKind::RecordEvent, + "turn_event" => QuotaUsageLinkKind::TurnEvent, + _ => QuotaUsageLinkKind::None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use statsai_core::{ + IdentitySource, LocationOrigin, QuotaCreditsV1, QuotaStatusV1, SourceAccountAssignment, + SourceAccountAssignmentId, SourceLocation, SOURCE_ACCOUNT_ASSIGNMENT_SCHEMA_VERSION, + }; + use std::collections::HashSet; + use std::path::Path; + + #[allow(clippy::too_many_arguments)] + fn sample_record( + source_id: SourceId, + observation_id: &str, + semantic_fingerprint: &str, + observed_at: DateTime, + reset_epoch: i64, + slot: &str, + window_minutes: u64, + used_percent: f64, + ) -> QuotaObservationRecordV1 { + let raw_rate_limits = serde_json::json!({ + slot: { + "window_minutes": window_minutes, + "used_percent": used_percent, + "resets_at": reset_epoch + }, + "credits": {"balance": "5.00"} + }); + let payload_hash = hash_text(&serde_json::to_string(&raw_rate_limits).expect("payload")); + QuotaObservationRecordV1 { + observation: QuotaObservationV1 { + schema_version: "quota_observation.v1".to_string(), + observation_id: observation_id.to_string(), + semantic_fingerprint: semantic_fingerprint.to_string(), + provider: "codex".to_string(), + source_id, + provider_account_id: None, + observed_at, + source_file_path_hash: format!("file-{observation_id}"), + source_record_id: format!("record-{observation_id}"), + source_line_number: 1, + payload_hash, + usage_sample: None, + usage_event_id: None, + usage_link_kind: QuotaUsageLinkKind::None, + status: QuotaStatusV1 { + plan_type: Some("pro".to_string()), + credits: QuotaCreditsV1 { + balance: Some("5".to_string()), + balance_raw: Some(serde_json::json!("5.00")), + ..QuotaCreditsV1::default() + }, + ..QuotaStatusV1::default() + }, + }, + windows: vec![QuotaWindowObservationV1 { + schema_version: "quota_window_observation.v1".to_string(), + window_observation_id: format!("window-{observation_id}"), + observation_id: observation_id.to_string(), + provider_slot: slot.to_string(), + limit_id: Some("subscription".to_string()), + window_minutes, + used_percent, + resets_at: DateTime::from_timestamp(reset_epoch, 0).expect("reset"), + resets_at_epoch_seconds: reset_epoch, + }], + raw_rate_limits, + } + } + + fn assigned_source(store: &Store, started_at: DateTime) -> (SourceId, ProviderAccountId) { + let source = SourceLocation::local_adapter( + "codex", + "codex-local-jsonl", + "test", + Path::new("/tmp/quota-source"), + LocationOrigin::Configured, + ); + store.upsert_source(&source).expect("source"); + let account_id = ProviderAccountId("account-codex".to_string()); + let now = Utc::now(); + store + .upsert_source_account_assignment(&SourceAccountAssignment { + schema_version: SOURCE_ACCOUNT_ASSIGNMENT_SCHEMA_VERSION.to_string(), + assignment_id: SourceAccountAssignmentId("assignment-quota".to_string()), + source_id: source.source_id.clone(), + provider: "codex".to_string(), + provider_account_id: account_id.clone(), + started_at, + ended_at: None, + record_source: IdentitySource::UserConfigured, + verified_at: Some(now), + created_at: now, + updated_at: now, + }) + .expect("assignment"); + (source.source_id, account_id) + } + + #[test] + fn semantic_collapse_prefers_attributed_linked_evidence() { + let now = Utc::now(); + let base = QuotaObservationV1 { + schema_version: "quota_observation.v1".to_string(), + observation_id: "one".to_string(), + semantic_fingerprint: "same".to_string(), + provider: "codex".to_string(), + source_id: SourceId("source-a".to_string()), + provider_account_id: None, + observed_at: now, + source_file_path_hash: "file".to_string(), + source_record_id: "record".to_string(), + source_line_number: 1, + payload_hash: "payload".to_string(), + usage_sample: None, + usage_event_id: None, + usage_link_kind: QuotaUsageLinkKind::None, + status: statsai_core::QuotaStatusV1::default(), + }; + let mut better = base.clone(); + better.observation_id = "two".to_string(); + better.provider_account_id = Some(ProviderAccountId("account".to_string())); + better.usage_event_id = Some(EventId("event".to_string())); + let records = collapse_semantic_duplicates(vec![ + QuotaObservationRecordV1 { + observation: base, + windows: Vec::new(), + raw_rate_limits: serde_json::json!({}), + }, + QuotaObservationRecordV1 { + observation: better, + windows: Vec::new(), + raw_rate_limits: serde_json::json!({}), + }, + ]); + assert_eq!(records.len(), 1); + assert_eq!(records[0].observation.observation_id, "two"); + } + + #[test] + fn quota_status_scopes_assignment_overlap_warnings_to_query() { + let store = Store::in_memory().expect("store"); + let started_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + for (source, provider, account, prefix) in [ + ("source-a", "codex", "account-a", "a"), + ("source-b", "codex", "account-b", "b"), + ("source-c", "claude", "account-c", "c"), + ] { + for (suffix, offset) in [("one", 0), ("two", 10)] { + store + .upsert_source_account_assignment(&SourceAccountAssignment { + schema_version: SOURCE_ACCOUNT_ASSIGNMENT_SCHEMA_VERSION.to_string(), + assignment_id: SourceAccountAssignmentId(format!( + "assignment-{prefix}-{suffix}" + )), + source_id: SourceId(source.to_string()), + provider: provider.to_string(), + provider_account_id: ProviderAccountId(account.to_string()), + started_at: started_at + Duration::seconds(offset), + ended_at: Some(started_at + Duration::seconds(offset + 20)), + record_source: IdentitySource::UserConfigured, + verified_at: Some(started_at), + created_at: started_at, + updated_at: started_at, + }) + .expect("assignment"); + } + } + + let account_status = store + .quota_status(&QuotaQuery { + provider_account_id: Some(ProviderAccountId("account-a".to_string())), + ..QuotaQuery::default() + }) + .expect("account status"); + assert_eq!( + account_status.assignment_overlap_warnings, + ["source source-a assignments assignment-a-one and assignment-a-two overlap"] + ); + + let provider_status = store + .quota_status(&QuotaQuery { + provider: Some("claude".to_string()), + ..QuotaQuery::default() + }) + .expect("provider status"); + assert_eq!( + provider_status.assignment_overlap_warnings, + ["source source-c assignments assignment-c-one and assignment-c-two overlap"] + ); + + let source_status = store + .quota_status(&QuotaQuery { + source_id: Some(SourceId("source-b".to_string())), + ..QuotaQuery::default() + }) + .expect("source status"); + assert_eq!( + source_status.assignment_overlap_warnings, + ["source source-b assignments assignment-b-one and assignment-b-two overlap"] + ); + + let later_status = store + .quota_status(&QuotaQuery { + from: Some(started_at + Duration::seconds(31)), + ..QuotaQuery::default() + }) + .expect("later status"); + assert!(later_status.assignment_overlap_warnings.is_empty()); + + let earlier_status = store + .quota_status(&QuotaQuery { + to: Some(started_at - Duration::seconds(1)), + ..QuotaQuery::default() + }) + .expect("earlier status"); + assert!(earlier_status.assignment_overlap_warnings.is_empty()); + } + + #[test] + fn quota_windows_keep_identical_observations_for_distinct_accounts() { + let store = Store::in_memory().expect("store"); + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + + for suffix in ["a", "b"] { + let source = SourceLocation::local_adapter( + "codex", + "codex-local-jsonl", + suffix, + Path::new(&format!("/tmp/quota-source-{suffix}")), + LocationOrigin::Configured, + ); + store.upsert_source(&source).expect("source"); + store + .upsert_source_account_assignment(&SourceAccountAssignment { + schema_version: SOURCE_ACCOUNT_ASSIGNMENT_SCHEMA_VERSION.to_string(), + assignment_id: SourceAccountAssignmentId(format!("assignment-{suffix}")), + source_id: source.source_id.clone(), + provider: "codex".to_string(), + provider_account_id: ProviderAccountId(format!("account-{suffix}")), + started_at: observed_at - Duration::seconds(1), + ended_at: None, + record_source: IdentitySource::UserConfigured, + verified_at: Some(observed_at), + created_at: observed_at, + updated_at: observed_at, + }) + .expect("assignment"); + let record = sample_record( + source.source_id, + &format!("observation-{suffix}"), + "same-semantic-fingerprint", + observed_at, + 1_787_500_000, + "primary", + 10_080, + 20.0, + ); + store + .upsert_quota_observations(&[record]) + .expect("quota observation"); + } + + let windows = store + .quota_windows(&QuotaQuery::default()) + .expect("quota windows"); + assert_eq!(windows.len(), 2); + assert_eq!( + windows + .iter() + .filter_map(|window| window.provider_account_id.as_ref()) + .collect::>() + .len(), + 2 + ); + } + + #[test] + fn quota_windows_keep_identical_unattributed_observations_in_distinct_source_scopes() { + let store = Store::in_memory().expect("store"); + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + + for suffix in ["a", "b"] { + let source = SourceLocation::local_adapter( + "codex", + "codex-local-jsonl", + suffix, + Path::new(&format!("/tmp/quota-unattributed-source-{suffix}")), + LocationOrigin::Configured, + ); + store.upsert_source(&source).expect("source"); + store + .upsert_quota_observations(&[sample_record( + source.source_id, + &format!("unattributed-{suffix}"), + "same-unattributed-semantic-fingerprint", + observed_at, + 1_787_500_000, + "primary", + 10_080, + 20.0, + )]) + .expect("quota observation"); + } + + let windows = store + .quota_windows(&QuotaQuery::default()) + .expect("quota windows"); + assert_eq!(windows.len(), 2); + assert_eq!( + windows + .iter() + .map(|window| window.window_id.as_str()) + .collect::>() + .len(), + 2 + ); + assert_eq!( + windows + .iter() + .filter_map(|window| window.source_id.as_ref()) + .collect::>() + .len(), + 2 + ); + assert!(windows.iter().all(|window| window.sample_count == 1)); + } + + #[test] + fn time_filtered_windows_do_not_query_usage_for_filtered_out_clusters() { + let store = Store::in_memory().expect("store"); + let old_observed_at = DateTime::from_timestamp(1_735_689_600, 0).expect("old observed"); + let old_reset = DateTime::from_timestamp(1_736_294_400, 0).expect("old reset"); + let recent_observed_at = + DateTime::from_timestamp(1_787_000_000, 0).expect("recent observed"); + let recent_reset = DateTime::from_timestamp(1_787_500_000, 0).expect("recent reset"); + let (source_id, account_id) = + assigned_source(&store, old_observed_at - Duration::seconds(1)); + store + .upsert_quota_observations(&[ + sample_record( + source_id.clone(), + "old-window", + "old-window", + old_observed_at, + old_reset.timestamp(), + "primary", + 10_080, + 10.0, + ), + sample_record( + source_id.clone(), + "recent-window", + "recent-window", + recent_observed_at, + recent_reset.timestamp(), + "primary", + 10_080, + 20.0, + ), + ]) + .expect("quota observations"); + store + .conn + .execute( + r#" + INSERT INTO usage_events ( + event_id, provider, source_id, provider_account_id, started_at, + total_tokens, semantic_fingerprint, payload + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + "#, + params![ + "invalid-old-event", + "codex", + &source_id.0, + &account_id.0, + (old_reset - Duration::days(1)).to_rfc3339(), + 0, + "invalid-old-event", + "not-json" + ], + ) + .expect("invalid legacy event"); + + let windows = store + .quota_windows(&QuotaQuery { + from: Some(recent_observed_at), + ..QuotaQuery::default() + }) + .expect("recent windows"); + + assert_eq!(windows.len(), 1); + assert_eq!(windows[0].first_observed_at, recent_observed_at); + + let mut limited = store + .quota_windows_without_usage_totals(&QuotaQuery::default()) + .expect("unenriched windows"); + assert_eq!(limited.len(), 2); + limited.truncate(1); + store + .enrich_quota_window_usage_totals(&mut limited) + .expect("enrich only the caller-selected recent window"); + assert_eq!(limited[0].first_observed_at, recent_observed_at); + } + + #[test] + fn quota_window_identity_and_evidence_are_stable_across_time_filters() { + let store = Store::in_memory().expect("store"); + let first_observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + let second_observed_at = first_observed_at + Duration::minutes(5); + let (source_id, _) = assigned_source(&store, first_observed_at - Duration::seconds(1)); + store + .upsert_quota_observations(&[ + sample_record( + source_id.clone(), + "stable-first", + "stable-first", + first_observed_at, + 1_787_500_000, + "primary", + 10_080, + 10.0, + ), + sample_record( + source_id, + "stable-second", + "stable-second", + second_observed_at, + 1_787_500_000, + "secondary", + 10_080, + 20.0, + ), + ]) + .expect("observations"); + + let full = store + .quota_windows(&QuotaQuery::default()) + .expect("full windows"); + let filtered = store + .quota_windows(&QuotaQuery { + from: Some(second_observed_at), + ..QuotaQuery::default() + }) + .expect("filtered windows"); + + assert_eq!(full.len(), 1); + assert_eq!(filtered.len(), 1); + assert_eq!(filtered[0].window_id, full[0].window_id); + assert_eq!(filtered[0].sample_count, 2); + assert_eq!(filtered[0].first_used_percent, 10.0); + assert_eq!(filtered[0].change_points, full[0].change_points); + } + + #[test] + fn semantic_collapse_drops_ambiguous_unattributed_copies() { + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + let mut unassigned = sample_record( + SourceId("source-unassigned".to_string()), + "ambiguous-unassigned", + "ambiguous-semantic", + observed_at, + 1_787_500_000, + "primary", + 10_080, + 20.0, + ); + let mut account_a = unassigned.clone(); + account_a.observation.observation_id = "ambiguous-a".to_string(); + account_a.observation.provider_account_id = + Some(ProviderAccountId("account-a".to_string())); + let mut account_b = unassigned.clone(); + account_b.observation.observation_id = "ambiguous-b".to_string(); + account_b.observation.provider_account_id = + Some(ProviderAccountId("account-b".to_string())); + unassigned.observation.provider_account_id = None; + + let collapsed = collapse_semantic_duplicates(vec![unassigned, account_a, account_b]); + + assert_eq!(collapsed.len(), 2); + assert!(collapsed + .iter() + .all(|record| record.observation.provider_account_id.is_some())); + } + + #[test] + fn two_device_projection_fixture_has_one_merge_scope_and_shared_point() { + let fixture = + include_str!("../../../docs/fixtures/quota_window_projection.v1.two_devices.json"); + let raw: serde_json::Value = serde_json::from_str(fixture).expect("fixture JSON"); + let projections: Vec = + serde_json::from_str(fixture).expect("projection fixture"); + + assert_eq!(projections.len(), 2); + assert_ne!(projections[0].device_id, projections[1].device_id); + assert_ne!(projections[0].projection_id, projections[1].projection_id); + assert_eq!(projections[0].provider, projections[1].provider); + assert_eq!( + projections[0].provider_account_id, + projections[1].provider_account_id + ); + assert_eq!(projections[0].limit_id, projections[1].limit_id); + assert_eq!(projections[0].window_minutes, projections[1].window_minutes); + assert!( + (projections[0].representative_reset_epoch_seconds + - projections[1].representative_reset_epoch_seconds) + .abs() + <= RESET_CLUSTER_TOLERANCE_SECONDS + ); + assert_eq!( + projections[0].change_points[0].point_fingerprint, + projections[1].change_points[0].point_fingerprint + ); + + for contribution in raw.as_array().expect("fixture array") { + for forbidden in [ + "source_id", + "usage_totals", + "total_tokens", + "estimated_cost_micro_usd", + "raw_rate_limits", + ] { + assert!(contribution.get(forbidden).is_none(), "found {forbidden}"); + } + } + } + + #[test] + fn quota_store_reuses_payloads_and_rescans_idempotently() { + let store = Store::in_memory().expect("store"); + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + let source_id = SourceId("source-a".to_string()); + let record = sample_record( + source_id, + "observation-a", + "semantic-a", + observed_at, + 1_787_500_000, + "primary", + 10_080, + 20.0, + ); + store + .upsert_quota_observations(std::slice::from_ref(&record)) + .expect("first upsert"); + store + .upsert_quota_observations(std::slice::from_ref(&record)) + .expect("rescan"); + + assert_eq!( + store + .quota_observations(&QuotaQuery::default(), false) + .expect("observations") + .len(), + 1 + ); + assert_eq!( + store + .conn + .query_row("SELECT COUNT(*) FROM quota_payloads", [], |row| { + row.get::<_, u64>(0) + }) + .expect("payload count"), + 1 + ); + } + + #[test] + fn quota_upsert_preserves_usage_link_only_for_matching_positive_sample() { + let store = Store::in_memory().expect("store"); + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + let mut linked = sample_record( + SourceId("source-link".to_string()), + "stable-line", + "semantic-original", + observed_at, + 1_787_500_000, + "primary", + 10_080, + 20.0, + ); + linked.observation.usage_sample = Some(statsai_core::UsageCounts { + input_tokens: Some(10), + total_tokens: Some(10), + ..statsai_core::UsageCounts::default() + }); + linked.observation.usage_event_id = Some(EventId("event-original".to_string())); + linked.observation.usage_link_kind = QuotaUsageLinkKind::RecordEvent; + store + .upsert_quota_observations(std::slice::from_ref(&linked)) + .expect("linked observation"); + + let mut unchanged = linked.clone(); + unchanged.observation.usage_event_id = None; + unchanged.observation.usage_link_kind = QuotaUsageLinkKind::None; + store + .upsert_quota_observations(std::slice::from_ref(&unchanged)) + .expect("unchanged archive observation"); + let stored = store + .quota_observations(&QuotaQuery::default(), false) + .expect("stored observation"); + assert_eq!( + stored[0].observation.usage_event_id, + Some(EventId("event-original".to_string())) + ); + + let mut changed = unchanged.clone(); + changed.observation.semantic_fingerprint = "semantic-changed".to_string(); + changed.observation.usage_sample = Some(statsai_core::UsageCounts { + input_tokens: Some(11), + total_tokens: Some(11), + ..statsai_core::UsageCounts::default() + }); + store + .upsert_quota_observations(std::slice::from_ref(&changed)) + .expect("changed archive observation"); + let stored = store + .quota_observations(&QuotaQuery::default(), false) + .expect("changed observation"); + assert_eq!(stored[0].observation.usage_event_id, None); + assert_eq!( + stored[0].observation.usage_link_kind, + QuotaUsageLinkKind::None + ); + + store + .upsert_quota_observations(std::slice::from_ref(&linked)) + .expect("restore linked observation"); + let mut missing = linked; + missing.observation.usage_sample = None; + missing.observation.usage_event_id = None; + missing.observation.usage_link_kind = QuotaUsageLinkKind::None; + store + .upsert_quota_observations(std::slice::from_ref(&missing)) + .expect("missing archive sample"); + let stored = store + .quota_observations(&QuotaQuery::default(), false) + .expect("missing observation"); + assert_eq!(stored[0].observation.usage_event_id, None); + } + + #[test] + fn quota_sync_projection_skips_local_usage_enrichment() { + let store = Store::in_memory().expect("store"); + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("observed"); + let reset = DateTime::from_timestamp(1_787_500_000, 0).expect("reset"); + let (source_id, account_id) = assigned_source(&store, observed_at - Duration::days(8)); + store + .upsert_quota_observations(&[sample_record( + source_id.clone(), + "projection-window", + "projection-window", + observed_at, + reset.timestamp(), + "primary", + 10_080, + 20.0, + )]) + .expect("quota observation"); + store + .conn + .execute( + r#" + INSERT INTO usage_events ( + event_id, provider, source_id, provider_account_id, started_at, + total_tokens, semantic_fingerprint, payload + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + "#, + params![ + "invalid-projection-event", + "codex", + &source_id.0, + &account_id.0, + (reset - Duration::days(1)).to_rfc3339(), + 0, + "invalid-projection-event", + "not-json" + ], + ) + .expect("invalid legacy event"); + + let projections = store + .quota_sync_projections(&QuotaQuery::default(), "device-a") + .expect("projection without local usage enrichment"); + + assert_eq!(projections.len(), 1); + } + + #[test] + fn quota_store_handles_ten_thousand_observations_with_one_repeated_payload() { + let store = Store::in_memory().expect("store"); + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + let base = sample_record( + SourceId("source-scale".to_string()), + "scale-0", + "scale-0", + observed_at, + 1_787_500_000, + "primary", + 10_080, + 20.0, + ); + let mut records = Vec::with_capacity(10_000); + for index in 0..10_000 { + let mut record = base.clone(); + record.observation.observation_id = format!("scale-{index}"); + record.observation.semantic_fingerprint = format!("semantic-scale-{index}"); + record.observation.source_record_id = format!("record-scale-{index}"); + record.observation.source_line_number = index as u64 + 1; + record.observation.observed_at = observed_at + Duration::seconds(index as i64); + record.windows[0].observation_id = record.observation.observation_id.clone(); + record.windows[0].window_observation_id = format!("window-scale-{index}"); + records.push(record); + } + store + .upsert_quota_observations(&records) + .expect("bulk observations"); + assert_eq!( + store + .quota_status(&QuotaQuery::default()) + .expect("status") + .total_observations, + 10_000 + ); + assert_eq!( + store + .conn + .query_row("SELECT COUNT(*) FROM quota_payloads", [], |row| { + row.get::<_, u64>(0) + }) + .expect("payload count"), + 1 + ); + } + + #[test] + fn quota_attribution_uses_exact_interval_boundary_and_reattributes_history() { + let store = Store::in_memory().expect("store"); + let boundary = DateTime::from_timestamp(1_787_000_000, 0).expect("boundary"); + let (source_id, account_id) = assigned_source(&store, boundary); + let before = sample_record( + source_id.clone(), + "before", + "before", + boundary - Duration::seconds(1), + 1_787_500_000, + "primary", + 10_080, + 10.0, + ); + let at = sample_record( + source_id.clone(), + "at", + "at", + boundary, + 1_787_500_100, + "secondary", + 10_080, + 11.0, + ); + store + .upsert_quota_observations(&[before, at]) + .expect("upsert"); + let observations = store + .quota_observations(&QuotaQuery::default(), false) + .expect("observations"); + assert!(observations[0].observation.provider_account_id.is_none()); + assert_eq!( + observations[1].observation.provider_account_id, + Some(account_id.clone()) + ); + let mut assignment = store + .list_source_account_assignments_for_source(&source_id) + .expect("assignments") + .remove(0); + assignment.started_at = boundary - Duration::minutes(1); + assignment.updated_at = Utc::now(); + store + .upsert_source_account_assignment(&assignment) + .expect("backdate assignment"); + store + .reattribute_quota_observations(&source_id) + .expect("reattribute"); + assert!(store + .quota_observations(&QuotaQuery::default(), false) + .expect("reattributed observations") + .iter() + .all(|record| record.observation.provider_account_id == Some(account_id.clone()))); + } + + #[test] + fn reconstruction_clusters_reset_drift_ignores_slot_and_projects_weekly_windows() { + let store = Store::in_memory().expect("store"); + let observed_at = DateTime::from_timestamp(1_787_000_000, 0).expect("time"); + let (source_id, account_id) = assigned_source(&store, observed_at - Duration::hours(1)); + let records = vec![ + sample_record( + source_id.clone(), + "one", + "one", + observed_at, + 1_787_500_000, + "primary", + 10_080, + 10.0, + ), + sample_record( + source_id, + "two", + "two", + observed_at + Duration::minutes(1), + 1_787_500_240, + "secondary", + 10_080, + 12.0, + ), + ]; + store.upsert_quota_observations(&records).expect("upsert"); + + let windows = store + .quota_windows(&QuotaQuery::default()) + .expect("windows"); + assert_eq!(windows.len(), 1); + assert_eq!(windows[0].sample_count, 2); + assert_eq!(windows[0].reset_min_epoch_seconds, 1_787_500_000); + assert_eq!(windows[0].reset_max_epoch_seconds, 1_787_500_240); + let projections = store + .quota_sync_projections(&QuotaQuery::default(), "device-a") + .expect("projections"); + let peer_projections = store + .quota_sync_projections(&QuotaQuery::default(), "device-b") + .expect("peer projections"); + assert_eq!(projections.len(), 1); + assert_eq!(peer_projections.len(), 1); + assert_eq!(projections[0].provider_account_id, account_id); + assert_ne!( + projections[0].projection_id, + peer_projections[0].projection_id + ); + assert_eq!( + projections[0].change_points[0].point_fingerprint, + peer_projections[0].change_points[0].point_fingerprint + ); + let projection_json = serde_json::to_value(&projections[0]).expect("json"); + assert!(projection_json.get("source_id").is_none()); + assert!(projection_json.get("total_tokens").is_none()); + assert!(projection_json.get("estimated_cost").is_none()); + } +} diff --git a/crates/statsai-store/tests/forward_schema.rs b/crates/statsai-store/tests/forward_schema.rs index 86d1f04..cbad9af 100644 --- a/crates/statsai-store/tests/forward_schema.rs +++ b/crates/statsai-store/tests/forward_schema.rs @@ -14,21 +14,21 @@ fn store_refuses_to_open_database_from_a_newer_binary() { applied_at TEXT NOT NULL ); INSERT INTO schema_migrations (version, applied_at) - VALUES (18, '2026-08-19T00:00:00Z'); + VALUES (19, '2026-08-19T00:00:00Z'); "#, ) .expect("record future schema"); drop(connection); let error = match Store::open(&path) { - Ok(_) => panic!("schema 17 binary must reject schema 18 database"), + Ok(_) => panic!("schema 18 binary must reject schema 19 database"), Err(error) => error, }; - assert_eq!(CURRENT_SCHEMA_VERSION, 17); + assert_eq!(CURRENT_SCHEMA_VERSION, 18); assert!(error .to_string() - .contains("database schema version 18 is newer than this StatsAI binary supports (17)")); + .contains("database schema version 19 is newer than this StatsAI binary supports (18)")); let connection = Connection::open(&path).expect("reopen rejected database"); assert_eq!( connection @@ -36,6 +36,6 @@ fn store_refuses_to_open_database_from_a_newer_binary() { row.get::<_, i64>(0) }) .expect("read future schema marker"), - 18 + 19 ); } diff --git a/crates/statsai/Cargo.toml b/crates/statsai/Cargo.toml index 60e096b..bad2b31 100644 --- a/crates/statsai/Cargo.toml +++ b/crates/statsai/Cargo.toml @@ -41,6 +41,7 @@ statsai-privacy.workspace = true anyhow.workspace = true chrono.workspace = true clap.workspace = true +csv.workspace = true schemars.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/crates/statsai/src/main.rs b/crates/statsai/src/main.rs index 790fe2e..53f0434 100644 --- a/crates/statsai/src/main.rs +++ b/crates/statsai/src/main.rs @@ -1,5 +1,5 @@ use anyhow::{bail, ensure, Context, Result}; -use chrono::{DateTime, Datelike, Duration, NaiveDate, Utc}; +use chrono::{DateTime, Datelike, Duration, Local, NaiveDate, Utc}; use clap::{Args, Parser, Subcommand}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -18,7 +18,8 @@ use statsai_core::{ sanitize_code_change_metric_for_sync, sanitize_task_bucket_for_sync, source_account_assignment_id, source_id as statsai_source_id, subscription_id, timestamp_in_period, ArchiveContentKind, ArchiveConversation, BillingPeriod, EventId, - IdentitySource, LocationOrigin, ProjectInfo, ProviderAccount, ProviderAccountId, ReportPeriod, + IdentitySource, LocationOrigin, ProjectInfo, ProviderAccount, ProviderAccountId, + QuotaObservationRecordV1, QuotaWindowSyncProjectionV1, QuotaWindowV1, ReportPeriod, SourceAccountAssignment, SourceAccountAssignmentId, SourceId, SourceKind, SourceLocation, SourceVerificationMode, Subscription, SubscriptionId, SubscriptionStatus, SyncAuthoritativeSnapshot, SyncBatch, TaskBucketSnapshot, TaskSpan, TaskStatus, TaskVerdict, @@ -41,7 +42,7 @@ use statsai_store::{apply_verified_source_state, verified_source_state_hash}; use statsai_store::{ close_active_verified_source_linkages, derive_task_work_items, find_existing_provider_account, reconcile_verified_source_state, upsert_provider_account, verified_source_observation_hash, - ScanFileStateEntry, Store, SyncPreferences, SyncState, TaskRebuildReport, + QuotaQuery, ScanFileStateEntry, Store, SyncPreferences, SyncState, TaskRebuildReport, UpsertProviderAccountInput, CURRENT_SCHEMA_VERSION, }; use statsai_sync::{ @@ -108,6 +109,8 @@ enum Command { Task(TaskCommand), #[command(about = "Collect and explore durable local conversation archives")] Conversation(ConversationCommand), + #[command(about = "Inspect reconstructed provider quota history")] + Quota(QuotaCommand), #[command(about = "Build and inspect the local privacy-filtered dataset")] Privacy(statsai::privacy_cli::PrivacyCommand), #[command(about = "Export a sync batch to a sink")] @@ -214,6 +217,67 @@ struct ConversationCommand { command: ConversationSubcommand, } +#[derive(Debug, Args)] +struct QuotaCommand { + #[command(subcommand)] + command: QuotaSubcommand, +} + +#[derive(Debug, Subcommand)] +enum QuotaSubcommand { + #[command(about = "Show quota collection and account-attribution coverage")] + Status { + #[arg(long, help = "Account id, email, provider id, or label")] + account: Option, + #[arg(long, help = "Output as JSON")] + json: bool, + }, + #[command(about = "Show the newest activated quota window")] + Current { + #[arg(long, help = "Account id, email, provider id, or label")] + account: Option, + #[arg(long, help = "Include shorter quota scopes")] + all_scopes: bool, + #[arg(long, help = "Include older epochs that overlap the selected window")] + include_overlaps: bool, + #[arg(long, help = "Output as JSON")] + json: bool, + }, + #[command(about = "List reconstructed quota windows")] + Windows { + #[arg(long, help = "Provider filter")] + provider: Option, + #[arg(long, help = "Account id, email, provider id, or label")] + account: Option, + #[arg(long, help = "Observation range start (date or RFC 3339)")] + from: Option, + #[arg(long, help = "Observation range end (date or RFC 3339)")] + to: Option, + #[arg(long, help = "Provider limit identity")] + limit_id: Option, + #[arg(long, default_value_t = 50, help = "Maximum windows to return")] + limit: usize, + #[arg(long, help = "Output as JSON")] + json: bool, + }, + #[command(about = "Show quota change points for one reconstructed window")] + History { + #[arg(long, help = "Reconstructed quota window id; newest when omitted")] + window_id: Option, + #[arg(long, help = "Include source observations and raw provider payloads")] + raw: bool, + #[arg(long, help = "Output as JSON")] + json: bool, + }, + #[command(about = "Export quota observations, windows, or weekly sync projections")] + Export { + #[arg(long, value_parser = ["observations", "windows", "sync-windows"])] + level: String, + #[arg(long, value_parser = ["csv", "json", "jsonl"])] + format: String, + }, +} + #[derive(Debug, Subcommand)] enum ConversationSubcommand { #[command(about = "Collect new or changed conversations from local provider sources")] @@ -860,6 +924,8 @@ enum StoreAdminSubcommand { enum SchemaSubcommand { #[command(about = "Print the sync_batch.v3 JSON Schema")] SyncBatch, + #[command(about = "Print the quota_window_sync_projection.v1 JSON Schema")] + QuotaWindowProjection, } #[derive(Debug, Args)] @@ -904,6 +970,7 @@ fn main() -> Result<()> { Command::Export(command) => export(command, &store), Command::Task(command) => task(command, &store), Command::Conversation(command) => conversation(command, &store, &device_id), + Command::Quota(command) => quota(command, &store, &device_id), Command::Privacy(command) => { statsai::privacy_cli::run(command, &store, &store_path) } @@ -989,6 +1056,7 @@ fn scan_with_adapters( let mut event_count = 0u64; let mut summary_count = 0u64; let mut task_span_count = 0u64; + let mut quota_observation_count = 0u64; let mut inserted_count = 0u64; let mut summary_written_count = 0u64; let mut task_span_written_count = 0u64; @@ -1137,11 +1205,13 @@ fn scan_with_adapters( let source_event_count = scan.events.len() as u64; let source_summary_count = scan.summaries.len() as u64; let source_task_span_count = scan.task_spans.len() as u64; + let source_quota_observation_count = scan.quota_observations.len() as u64; let has_scan_activity = touched_files || (has_cache_entry_upgrades && !command.preview) || source_event_count > 0 || source_summary_count > 0 || source_task_span_count > 0 + || source_quota_observation_count > 0 || scan.diagnostics.files_scanned > 0 || scan.diagnostics.files_skipped_unchanged > 0 || log_rows > 0 @@ -1151,6 +1221,7 @@ fn scan_with_adapters( && source_event_count == 0 && source_summary_count == 0 && source_task_span_count == 0 + && source_quota_observation_count == 0 && !touched_files && !has_cache_entry_upgrades && !verified_state_changed; @@ -1164,6 +1235,7 @@ fn scan_with_adapters( event_count += source_event_count; summary_count += source_summary_count; task_span_count += source_task_span_count; + quota_observation_count += source_quota_observation_count; total_usage.add_totals(&source_usage); total_summary_usage.add_totals(&source_summary_usage); add_diagnostics(&mut total_diagnostics, &scan.diagnostics); @@ -1194,6 +1266,7 @@ fn scan_with_adapters( usage: &source_usage, summaries: source_summary_count, task_spans: source_task_span_count, + quota_observations: source_quota_observation_count, summary_usage: &source_summary_usage, diagnostics: &scan.diagnostics, verbose: command.verbose || command.explain, @@ -1229,6 +1302,9 @@ fn scan_with_adapters( store.delete_events_for_sources(std::slice::from_ref(&source.source_id))?; removed_summary_count += store .delete_summaries_for_sources(std::slice::from_ref(&source.source_id))?; + store.delete_quota_observations_for_sources(std::slice::from_ref( + &source.source_id, + ))?; if command.include_tasks { let deleted = store.delete_task_spans_for_sources(std::slice::from_ref( &source.source_id, @@ -1253,6 +1329,10 @@ fn scan_with_adapters( &source.source_id, &reconciled_file_hashes, )?; + store.delete_quota_observations_for_source_file_hashes( + &source.source_id, + &reconciled_file_hashes, + )?; if command.include_tasks { let deleted = store.delete_task_spans_for_source_file_hashes( &source.source_id, @@ -1269,6 +1349,12 @@ fn scan_with_adapters( let insert_result = store.insert_events_with_resolution(&scan.events)?; inserted_count += insert_result.inserted; insert_events_duration_ms += insert_started_at.elapsed().as_millis() as u64; + rewrite_quota_usage_event_ids( + &mut scan.quota_observations, + &insert_result.canonical_event_ids, + ); + store.upsert_quota_observations(&scan.quota_observations)?; + store.clear_orphaned_quota_usage_links()?; if command.include_tasks { rewrite_task_span_linked_event_ids( &mut scan.task_spans, @@ -1342,10 +1428,11 @@ fn scan_with_adapters( if command.preview { if command.verbose { println!( - "preview total: sources={} usage_events={} summaries={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} log_rows={} written=0", + "preview total: sources={} usage_events={} summaries={} quota_observations={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} log_rows={} written=0", format_u64(total_sources), format_u64(event_count), format_u64(summary_count), + format_u64(quota_observation_count), format_u64(total_usage.input_tokens), format_u64(total_usage.cache_creation_tokens), format_u64(total_usage.cached_input_tokens), @@ -1370,10 +1457,11 @@ fn scan_with_adapters( print_scan_diagnostics_total(&total_diagnostics); } else { println!( - "preview total: sources={} usage_events={} summaries={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} written=0", + "preview total: sources={} usage_events={} summaries={} quota_observations={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} written=0", format_u64(total_sources), format_u64(event_count), format_u64(summary_count), + format_u64(quota_observation_count), format_u64(total_usage.input_tokens), format_u64(total_usage.cache_creation_tokens), format_u64(total_usage.cached_input_tokens), @@ -1391,7 +1479,7 @@ fn scan_with_adapters( } } else { println!( - "scan complete: sources={} usage_events={} inserted={} summaries={} summaries_written={} task_spans={} task_spans_written={} work_items_rebuilt={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} log_rows={}", + "scan complete: sources={} usage_events={} inserted={} summaries={} summaries_written={} task_spans={} task_spans_written={} quota_observations={} work_items_rebuilt={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} log_rows={}", format_u64(total_sources), format_u64(event_count), format_u64(inserted_count), @@ -1399,6 +1487,7 @@ fn scan_with_adapters( format_u64(summary_written_count), format_u64(task_span_count), format_u64(task_span_written_count), + format_u64(quota_observation_count), format_u64(rebuilt_work_item_count), format_u64(total_usage.input_tokens), format_u64(total_usage.cache_creation_tokens), @@ -1647,6 +1736,14 @@ fn source(command: SourceCommand, store: &Store, device_id: &str) -> Result<()> } else { Default::default() }; + let deleted_quota_observations = if delete_data { + store.delete_quota_observations_for_sources(std::slice::from_ref(&source_id))? + } else { + 0 + }; + if delete_data { + store.clear_orphaned_quota_usage_links()?; + } let deleted_trace_edits = if delete_data { store.delete_archive_import_for_sources(std::slice::from_ref(&source_id))? } else { @@ -1683,6 +1780,7 @@ fn source(command: SourceCommand, store: &Store, device_id: &str) -> Result<()> "deleted_summaries": deleted_summaries, "deleted_scan_cache_entries": deleted_scan_entries, "deleted_task_spans": deleted_task_spans.deleted, + "deleted_quota_observations": deleted_quota_observations, "deleted_code_change_traces": deleted_trace_edits, "work_items_rebuilt": rebuilt_work_items, "code_change_metrics_rebuilt": rebuilt_code_change_metrics, @@ -2490,11 +2588,16 @@ struct AccountReferenceCounts { subscriptions: usize, events: usize, summaries: usize, + quota_observations: usize, } impl AccountReferenceCounts { fn total(&self) -> usize { - self.source_account_assignments + self.subscriptions + self.events + self.summaries + self.source_account_assignments + + self.subscriptions + + self.events + + self.summaries + + self.quota_observations } } @@ -2629,12 +2732,13 @@ fn remove_orphan_provider_account( account_reference_counts(store, &account.provider_account_id, Some(provider))?; if remaining_references.total() > 0 { bail!( - "account {} still has references: {} source assignments, {} subscriptions, {} events, {} summaries", + "account {} still has references: {} source assignments, {} subscriptions, {} events, {} summaries, {} quota observations", display_account_identity(&account), remaining_references.source_account_assignments, remaining_references.subscriptions, remaining_references.events, - remaining_references.summaries + remaining_references.summaries, + remaining_references.quota_observations ); } let deleted = if dry_run { @@ -2831,12 +2935,21 @@ fn account_reference_counts( .filter(|summary| summary.provider_account_id.as_ref() == Some(provider_account_id)) .filter(|summary| provider_matches(&summary.provider)) .count(); + let quota_observations = store + .quota_observations(&QuotaQuery::default(), false)? + .into_iter() + .filter(|record| { + record.observation.provider_account_id.as_ref() == Some(provider_account_id) + }) + .filter(|record| provider_matches(&record.observation.provider)) + .count(); Ok(AccountReferenceCounts { source_account_assignments, subscriptions, events, summaries, + quota_observations, }) } @@ -5952,6 +6065,663 @@ fn schema(command: SchemaCommand) -> Result<()> { let schema = schemars::schema_for!(SyncBatch); println!("{}", serde_json::to_string_pretty(&schema)?); } + SchemaSubcommand::QuotaWindowProjection => { + let schema = schemars::schema_for!(QuotaWindowSyncProjectionV1); + println!("{}", serde_json::to_string_pretty(&schema)?); + } + } + Ok(()) +} + +fn quota(command: QuotaCommand, store: &Store, device_id: &str) -> Result<()> { + match command.command { + QuotaSubcommand::Status { account, json } => { + let query = quota_query(store, None, account.as_deref(), None, None, None)?; + let status = store.quota_status(&query)?; + if json { + println!("{}", serde_json::to_string_pretty(&status)?); + } else { + println!( + "quota observations: {}", + format_u64(status.total_observations) + ); + println!( + "distinct: {} ({} copied duplicates collapsed)", + format_u64(status.distinct_observations), + format_u64(status.duplicate_observations) + ); + println!( + "attributed: {} range: {}", + format_u64(status.attributed_observations), + format_quota_range(status.attributed_range.as_ref()) + ); + println!( + "unattributed: {} range: {}", + format_u64(status.unattributed_observations), + format_quota_range(status.unattributed_range.as_ref()) + ); + println!( + "weekly sync coverage: {}/{} ({:.1}%)", + format_u64(status.weekly_sync_eligible_observations), + format_u64(status.weekly_observations), + status.weekly_sync_eligible_coverage_percent + ); + if status.unattributed_observations > 0 { + eprintln!( + "warning: historical quota observations remain unassigned; backdate a source connection with `statsai source connect --started-at ...`" + ); + } + for warning in status.assignment_overlap_warnings { + eprintln!("warning: {warning}"); + } + } + } + QuotaSubcommand::Current { + account, + all_scopes, + include_overlaps, + json, + } => { + let query = quota_query(store, None, account.as_deref(), None, None, None)?; + let mut windows = select_current_quota_windows( + store.quota_windows_without_usage_totals(&query)?, + all_scopes, + include_overlaps, + ); + store.enrich_quota_window_usage_totals(&mut windows)?; + print_quota_windows(&windows, json)?; + } + QuotaSubcommand::Windows { + provider, + account, + from, + to, + limit_id, + limit, + json, + } => { + let query = quota_query( + store, + provider.as_deref(), + account.as_deref(), + from.as_deref(), + to.as_deref(), + limit_id, + )?; + let mut windows = store.quota_windows_without_usage_totals(&query)?; + windows.truncate(limit.min(10_000)); + store.enrich_quota_window_usage_totals(&mut windows)?; + print_quota_windows(&windows, json)?; + } + QuotaSubcommand::History { + window_id, + raw, + json, + } => { + let query = QuotaQuery::default(); + let windows = store.quota_windows_without_usage_totals(&query)?; + let window = window_id + .as_deref() + .map(|id| windows.iter().find(|window| window.window_id == id)) + .unwrap_or_else(|| windows.first()) + .with_context(|| { + window_id.map_or_else( + || "no quota windows found".to_string(), + |id| format!("quota window not found: {id}"), + ) + })?; + if raw { + let observations = + raw_observations_for_window(store.quota_observations(&query, false)?, window); + if json { + println!( + "{}", + serde_json::to_string_pretty(&json!({ + "schema_version": "quota_history.v1", + "window_id": window.window_id, + "raw": true, + "observations": observations, + }))? + ); + } else { + for record in observations { + println!("{}", serde_json::to_string_pretty(&record)?); + } + } + } else if json { + println!( + "{}", + serde_json::to_string_pretty(&json!({ + "schema_version": "quota_history.v1", + "window_id": window.window_id, + "raw": false, + "change_points": window.change_points, + }))? + ); + } else { + println!("{}", quota_window_heading(window)); + for point in &window.change_points { + println!( + " {} {:>6.2}% reset {} slot={}", + format_local_timestamp(point.observed_at), + point.used_percent, + format_local_timestamp(point.resets_at), + point.provider_slot, + ); + } + } + } + QuotaSubcommand::Export { level, format } => { + let query = QuotaQuery::default(); + match level.as_str() { + "observations" => { + export_quota_observations(&store.quota_observations(&query, false)?, &format)? + } + "windows" => export_quota_windows(&store.quota_windows(&query)?, &format)?, + "sync-windows" => { + let status = store.quota_status(&query)?; + if status.weekly_sync_eligible_observations < status.weekly_observations { + eprintln!( + "warning: only {}/{} weekly observations are attributed and eligible for sync projection export", + status.weekly_sync_eligible_observations, + status.weekly_observations + ); + } + export_quota_projections( + &store.quota_sync_projections(&query, device_id)?, + &format, + )?; + } + _ => unreachable!("clap validates quota export level"), + } + } + } + Ok(()) +} + +fn quota_query( + store: &Store, + provider: Option<&str>, + account: Option<&str>, + from: Option<&str>, + to: Option<&str>, + limit_id: Option, +) -> Result { + let provider = provider.map(canonical_provider).transpose()?; + let provider_account_id = account + .map(|selector| resolve_quota_account_selector(store, provider.as_deref(), selector)) + .transpose()? + .map(|account| account.provider_account_id); + Ok(QuotaQuery { + provider, + provider_account_id, + source_id: None, + from: from.map(parse_date).transpose()?, + to: to.map(parse_quota_range_end).transpose()?, + limit_id, + }) +} + +fn parse_quota_range_end(value: &str) -> Result> { + if let Ok(timestamp) = DateTime::parse_from_rfc3339(value) { + return Ok(timestamp.with_timezone(&Utc)); + } + let date = NaiveDate::parse_from_str(value, "%Y-%m-%d")?; + let next_day = date + .succ_opt() + .context("quota range end is outside the supported calendar")?; + Ok(next_day + .and_hms_opt(0, 0, 0) + .context("failed to build quota range end")? + .and_utc() + - Duration::nanoseconds(1)) +} + +fn resolve_quota_account_selector( + store: &Store, + provider: Option<&str>, + selector: &str, +) -> Result { + let selector = selector.trim(); + if selector.is_empty() { + bail!("account selector cannot be empty"); + } + let normalized_email = normalize_email(selector); + let normalized_provider_user_id = normalize_provider_user_id(selector); + let normalized_label = selector.to_ascii_lowercase(); + let mut matches = store + .list_accounts()? + .into_iter() + .filter(|account| provider.is_none_or(|provider| account.provider == provider)) + .filter(|account| { + account.provider_account_id.0 == selector + || account.email.as_deref().map(normalize_email).as_deref() + == Some(normalized_email.as_str()) + || account + .provider_user_id + .as_deref() + .map(normalize_provider_user_id) + .as_deref() + == Some(normalized_provider_user_id.as_str()) + || account + .account_label + .as_deref() + .map(|label| label.trim().to_ascii_lowercase()) + .as_deref() + == Some(normalized_label.as_str()) + }) + .collect::>(); + matches.sort_by_key(|account| account.provider_account_id.0.clone()); + matches.dedup_by(|left, right| left.provider_account_id == right.provider_account_id); + match matches.len() { + 0 => bail!("no account matched '{selector}'"), + 1 => Ok(matches.remove(0)), + _ => bail!( + "multiple provider accounts matched '{selector}'; use a stable provider account id" + ), + } +} + +fn select_current_quota_windows( + windows: Vec, + all_scopes: bool, + include_overlaps: bool, +) -> Vec { + let mut by_scope = BTreeMap::< + (String, Option, Option, Option, u64), + Vec, + >::new(); + for window in windows { + by_scope + .entry(( + window.provider.clone(), + window.provider_account_id.as_ref().map(|id| id.0.clone()), + window.source_id.as_ref().map(|id| id.0.clone()), + window.limit_id.clone(), + window.window_minutes, + )) + .or_default() + .push(window); + } + let mut selected = Vec::new(); + for mut scope_windows in by_scope.into_values() { + scope_windows.sort_by_key(|window| window.first_observed_at); + let newest = scope_windows.pop().expect("scope contains a window"); + if include_overlaps { + selected.extend(scope_windows.into_iter().filter(|older| { + older.inferred_start < newest.representative_reset + && older.representative_reset > newest.inferred_start + })); + } + selected.push(newest); + } + if !all_scopes { + let mut longest = + BTreeMap::<(String, Option, Option, Option), u64>::new(); + for window in &selected { + let key = ( + window.provider.clone(), + window.provider_account_id.as_ref().map(|id| id.0.clone()), + window.source_id.as_ref().map(|id| id.0.clone()), + window.limit_id.clone(), + ); + longest + .entry(key) + .and_modify(|duration| *duration = (*duration).max(window.window_minutes)) + .or_insert(window.window_minutes); + } + selected.retain(|window| { + longest.get(&( + window.provider.clone(), + window.provider_account_id.as_ref().map(|id| id.0.clone()), + window.source_id.as_ref().map(|id| id.0.clone()), + window.limit_id.clone(), + )) == Some(&window.window_minutes) + }); + } + selected.sort_by(|left, right| { + right + .window_minutes + .cmp(&left.window_minutes) + .then_with(|| right.first_observed_at.cmp(&left.first_observed_at)) + }); + selected +} + +fn print_quota_windows(windows: &[QuotaWindowV1], json_output: bool) -> Result<()> { + if json_output { + println!("{}", serde_json::to_string_pretty(windows)?); + return Ok(()); + } + if windows.is_empty() { + println!("no quota windows"); + return Ok(()); + } + for window in windows { + println!("{}", quota_window_heading(window)); + println!( + " observed {} to {}; samples={} transition={:?}", + format_local_timestamp(window.first_observed_at), + format_local_timestamp(window.last_observed_at), + format_u64(window.sample_count), + window.transition, + ); + if let Some(usage_totals) = &window.usage_totals { + println!( + " usage events={} tokens={} cost_micro_usd={}", + format_u64(usage_totals.event_count), + format_u64(usage_totals.total_tokens), + usage_totals + .estimated_cost_micro_usd + .map_or_else(|| "unknown".to_string(), |value| value.to_string()), + ); + } else { + println!(" usage unavailable until the quota history is attributed to an account"); + } + if window.has_schedule_overlap { + eprintln!( + "warning: {} overlaps another reconstructed epoch in the same quota scope", + window.window_id + ); + } + } + Ok(()) +} + +fn quota_window_heading(window: &QuotaWindowV1) -> String { + let identity = window.provider_account_id.as_ref().map_or_else( + || { + window.source_id.as_ref().map_or_else( + || "unassigned".to_string(), + |id| format!("unassigned@{}", id.0), + ) + }, + |id| id.0.clone(), + ); + format!( + "{} {} {}m {:.2}% reset {} id={}", + window.provider, + identity, + window.window_minutes, + window.latest_used_percent, + format_local_timestamp(window.representative_reset), + window.window_id, + ) +} + +fn format_local_timestamp(timestamp: DateTime) -> String { + timestamp + .with_timezone(&Local) + .format("%Y-%m-%d %H:%M:%S %Z") + .to_string() +} + +fn format_quota_range(range: Option<&statsai_store::QuotaDateRange>) -> String { + range.map_or_else( + || "none".to_string(), + |range| { + format!( + "{} to {}", + format_local_timestamp(range.first), + format_local_timestamp(range.last) + ) + }, + ) +} + +fn raw_observations_for_window( + records: Vec, + window: &QuotaWindowV1, +) -> Vec { + records + .into_iter() + .filter(|record| { + record.observation.provider == window.provider + && record.observation.provider_account_id == window.provider_account_id + && window + .source_id + .as_ref() + .is_none_or(|source_id| &record.observation.source_id == source_id) + && record.windows.iter().any(|candidate| { + candidate.limit_id == window.limit_id + && candidate.window_minutes == window.window_minutes + && candidate.resets_at_epoch_seconds >= window.reset_min_epoch_seconds + && candidate.resets_at_epoch_seconds <= window.reset_max_epoch_seconds + }) + }) + .collect() +} + +fn export_quota_observations(records: &[QuotaObservationRecordV1], format: &str) -> Result<()> { + match format { + "json" => println!("{}", serde_json::to_string_pretty(records)?), + "jsonl" => print_json_lines(records)?, + "csv" => { + let mut writer = csv::Writer::from_writer(std::io::stdout()); + writer.write_record([ + "schema_version", + "observation_id", + "semantic_fingerprint", + "provider", + "source_id", + "provider_account_id", + "observed_at", + "source_file_path_hash", + "source_record_id", + "usage_event_id", + "usage_link_kind", + "payload_hash", + "usage_total_tokens", + "provider_slot", + "limit_id", + "window_minutes", + "used_percent", + "resets_at", + "resets_at_epoch_seconds", + ])?; + for record in records { + let observation = &record.observation; + let base = vec![ + observation.schema_version.clone(), + observation.observation_id.clone(), + observation.semantic_fingerprint.clone(), + observation.provider.clone(), + observation.source_id.0.clone(), + observation + .provider_account_id + .as_ref() + .map(|id| id.0.clone()) + .unwrap_or_default(), + observation.observed_at.to_rfc3339(), + observation.source_file_path_hash.clone(), + observation.source_record_id.clone(), + observation + .usage_event_id + .as_ref() + .map(|id| id.0.clone()) + .unwrap_or_default(), + serde_json::to_value(observation.usage_link_kind)? + .as_str() + .unwrap_or("none") + .to_string(), + observation.payload_hash.clone(), + observation + .usage_sample + .as_ref() + .map(|usage| usage.computed_total().to_string()) + .unwrap_or_default(), + ]; + if record.windows.is_empty() { + let mut row = base; + row.extend(std::iter::repeat_n(String::new(), 6)); + writer.write_record(row)?; + continue; + } + for window in &record.windows { + let mut row = base.clone(); + row.extend([ + window.provider_slot.clone(), + window.limit_id.clone().unwrap_or_default(), + window.window_minutes.to_string(), + window.used_percent.to_string(), + window.resets_at.to_rfc3339(), + window.resets_at_epoch_seconds.to_string(), + ]); + writer.write_record(row)?; + } + } + writer.flush()?; + } + _ => unreachable!("clap validates quota export format"), + } + Ok(()) +} + +fn export_quota_windows(windows: &[QuotaWindowV1], format: &str) -> Result<()> { + match format { + "json" => println!("{}", serde_json::to_string_pretty(windows)?), + "jsonl" => print_json_lines(windows)?, + "csv" => { + let mut writer = csv::Writer::from_writer(std::io::stdout()); + writer.write_record([ + "schema_version", + "window_id", + "provider", + "provider_account_id", + "source_id", + "limit_id", + "window_minutes", + "inferred_start", + "representative_reset", + "representative_reset_epoch_seconds", + "reset_min_epoch_seconds", + "reset_max_epoch_seconds", + "first_observed_at", + "last_observed_at", + "sample_count", + "first_used_percent", + "latest_used_percent", + "minimum_used_percent", + "maximum_used_percent", + "event_count", + "total_tokens", + "estimated_cost_micro_usd", + ])?; + for window in windows { + writer.write_record([ + window.schema_version.clone(), + window.window_id.clone(), + window.provider.clone(), + window + .provider_account_id + .as_ref() + .map(|id| id.0.clone()) + .unwrap_or_default(), + window + .source_id + .as_ref() + .map(|id| id.0.clone()) + .unwrap_or_default(), + window.limit_id.clone().unwrap_or_default(), + window.window_minutes.to_string(), + window.inferred_start.to_rfc3339(), + window.representative_reset.to_rfc3339(), + window.representative_reset_epoch_seconds.to_string(), + window.reset_min_epoch_seconds.to_string(), + window.reset_max_epoch_seconds.to_string(), + window.first_observed_at.to_rfc3339(), + window.last_observed_at.to_rfc3339(), + window.sample_count.to_string(), + window.first_used_percent.to_string(), + window.latest_used_percent.to_string(), + window.minimum_used_percent.to_string(), + window.maximum_used_percent.to_string(), + window + .usage_totals + .as_ref() + .map(|totals| totals.event_count.to_string()) + .unwrap_or_default(), + window + .usage_totals + .as_ref() + .map(|totals| totals.total_tokens.to_string()) + .unwrap_or_default(), + window + .usage_totals + .as_ref() + .and_then(|totals| totals.estimated_cost_micro_usd) + .map(|cost| cost.to_string()) + .unwrap_or_default(), + ])?; + } + writer.flush()?; + } + _ => unreachable!("clap validates quota export format"), + } + Ok(()) +} + +fn export_quota_projections( + projections: &[QuotaWindowSyncProjectionV1], + format: &str, +) -> Result<()> { + match format { + "json" => println!("{}", serde_json::to_string_pretty(projections)?), + "jsonl" => print_json_lines(projections)?, + "csv" => { + let mut writer = csv::Writer::from_writer(std::io::stdout()); + writer.write_record([ + "schema_version", + "projection_id", + "device_id", + "provider", + "provider_account_id", + "limit_id", + "window_minutes", + "inferred_start", + "representative_reset", + "representative_reset_epoch_seconds", + "reset_min_epoch_seconds", + "reset_max_epoch_seconds", + "first_observed_at", + "last_observed_at", + "sample_count", + "latest_used_percent", + "change_points_json", + "status_json", + ])?; + for projection in projections { + writer.write_record([ + projection.schema_version.clone(), + projection.projection_id.clone(), + projection.device_id.clone(), + projection.provider.clone(), + projection.provider_account_id.0.clone(), + projection.limit_id.clone().unwrap_or_default(), + projection.window_minutes.to_string(), + projection.inferred_start.to_rfc3339(), + projection.representative_reset.to_rfc3339(), + projection.representative_reset_epoch_seconds.to_string(), + projection.reset_min_epoch_seconds.to_string(), + projection.reset_max_epoch_seconds.to_string(), + projection.first_observed_at.to_rfc3339(), + projection.last_observed_at.to_rfc3339(), + projection.sample_count.to_string(), + projection.latest_used_percent.to_string(), + serde_json::to_string(&projection.change_points)?, + serde_json::to_string(&projection.latest_status)?, + ])?; + } + writer.flush()?; + } + _ => unreachable!("clap validates quota export format"), + } + Ok(()) +} + +fn print_json_lines(values: &[T]) -> Result<()> { + for value in values { + println!("{}", serde_json::to_string(value)?); } Ok(()) } @@ -6302,6 +7072,7 @@ fn collect_archive_source_entries( &scan.artifact_dependencies, &scan.trace_edits, scan.trace_coverage, + &scan.quota_observations, )?; collected.files += scan.diagnostics.files_scanned; collected.conversations += write.conversations; @@ -7777,6 +8548,7 @@ fn reattribute_source_records(store: &Store, source_id: &SourceId) -> Result<()> } store.rewrite_events(&events)?; store.rewrite_summaries(&summaries)?; + store.reattribute_quota_observations(source_id)?; Ok(()) } @@ -8004,6 +8776,7 @@ struct ScanPreviewLine<'a> { usage: &'a UsageTotals, summaries: u64, task_spans: u64, + quota_observations: u64, summary_usage: &'a UsageTotals, diagnostics: &'a ScanDiagnostics, verbose: bool, @@ -8012,12 +8785,13 @@ struct ScanPreviewLine<'a> { fn print_scan_preview_line(line: ScanPreviewLine<'_>) { if line.verbose { println!( - "{} path={} usage_events={} summaries={} task_spans={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} raw_rows={} candidates={} duplicates={} skipped_zero={} invalid={} files={} cached={} timestamp_fallbacks={} model_fallbacks={} origin={} source={}", + "{} path={} usage_events={} summaries={} task_spans={} quota_observations={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={} raw_rows={} candidates={} duplicates={} skipped_zero={} invalid={} files={} cached={} timestamp_fallbacks={} model_fallbacks={} origin={} source={}", line.source.provider, preview_path_label(line.source), line.usage_events, line.summaries, line.task_spans, + line.quota_observations, format_u64(line.usage.input_tokens), format_u64(line.usage.cache_creation_tokens), format_u64(line.usage.cached_input_tokens), @@ -8040,12 +8814,13 @@ fn print_scan_preview_line(line: ScanPreviewLine<'_>) { ); } else { println!( - "{} path={} usage_events={} summaries={} task_spans={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={}", + "{} path={} usage_events={} summaries={} task_spans={} quota_observations={} input={} cache_create={} cache_read={} output={} total={} est_cost={} summary_total={} summary_est_cost={}", line.source.provider, preview_path_label(line.source), line.usage_events, line.summaries, line.task_spans, + line.quota_observations, format_u64(line.usage.input_tokens), format_u64(line.usage.cache_creation_tokens), format_u64(line.usage.cached_input_tokens), @@ -8255,6 +9030,20 @@ fn rewrite_task_span_linked_event_ids( } } +fn rewrite_quota_usage_event_ids( + observations: &mut [QuotaObservationRecordV1], + canonical_event_ids: &HashMap, +) { + for record in observations { + let Some(event_id) = record.observation.usage_event_id.as_ref() else { + continue; + }; + if let Some(canonical) = canonical_event_ids.get(event_id) { + record.observation.usage_event_id = Some(canonical.clone()); + } + } +} + #[derive(Debug, Clone, Copy, Default)] struct TaskSpanRuntimeRollup { total_messages: u64, @@ -12240,6 +13029,7 @@ mod tests { &[], &edits, statsai_core::CoverageStatus::Complete, + &[], ) .expect("seed trace edits"); } @@ -20180,4 +20970,116 @@ mod tests { assert_eq!(value["price"], json!(20.0)); assert_eq!(value["price_cents"], json!(2000)); } + + fn test_unattributed_quota_window(source_id: &str, window_id: &str) -> QuotaWindowV1 { + let observed_at = Utc + .with_ymd_and_hms(2026, 8, 20, 12, 0, 0) + .single() + .expect("observed at"); + let reset = observed_at + Duration::days(7); + QuotaWindowV1 { + schema_version: "quota_window.v1".to_string(), + window_id: window_id.to_string(), + provider: "codex".to_string(), + provider_account_id: None, + source_id: Some(SourceId(source_id.to_string())), + limit_id: Some("subscription".to_string()), + window_minutes: 10_080, + inferred_start: reset - Duration::days(7), + representative_reset: reset, + representative_reset_epoch_seconds: reset.timestamp(), + reset_min: reset, + reset_min_epoch_seconds: reset.timestamp(), + reset_max: reset, + reset_max_epoch_seconds: reset.timestamp(), + first_observed_at: observed_at, + last_observed_at: observed_at, + sample_count: 1, + first_used_percent: 20.0, + latest_used_percent: 20.0, + minimum_used_percent: 20.0, + maximum_used_percent: 20.0, + transition: statsai_core::QuotaTransitionKind::Initial, + has_schedule_overlap: false, + change_points: Vec::new(), + latest_status: statsai_core::QuotaStatusV1::default(), + usage_totals: None, + } + } + + fn test_unattributed_quota_record(source_id: &str) -> QuotaObservationRecordV1 { + let observed_at = Utc + .with_ymd_and_hms(2026, 8, 20, 12, 0, 0) + .single() + .expect("observed at"); + let reset = observed_at + Duration::days(7); + QuotaObservationRecordV1 { + observation: statsai_core::QuotaObservationV1 { + schema_version: "quota_observation.v1".to_string(), + observation_id: format!("observation-{source_id}"), + semantic_fingerprint: format!("semantic-{source_id}"), + provider: "codex".to_string(), + source_id: SourceId(source_id.to_string()), + provider_account_id: None, + observed_at, + source_file_path_hash: format!("file-{source_id}"), + source_record_id: format!("record-{source_id}"), + source_line_number: 1, + payload_hash: format!("payload-{source_id}"), + usage_sample: None, + usage_event_id: None, + usage_link_kind: statsai_core::QuotaUsageLinkKind::None, + status: statsai_core::QuotaStatusV1::default(), + }, + windows: vec![statsai_core::QuotaWindowObservationV1 { + schema_version: "quota_window_observation.v1".to_string(), + window_observation_id: format!("window-observation-{source_id}"), + observation_id: format!("observation-{source_id}"), + provider_slot: "primary".to_string(), + limit_id: Some("subscription".to_string()), + window_minutes: 10_080, + used_percent: 20.0, + resets_at: reset, + resets_at_epoch_seconds: reset.timestamp(), + }], + raw_rate_limits: json!({}), + } + } + + #[test] + fn current_quota_windows_keep_unattributed_source_scopes_separate() { + let selected = select_current_quota_windows( + vec![ + test_unattributed_quota_window("source-a", "window-a"), + test_unattributed_quota_window("source-b", "window-b"), + ], + false, + false, + ); + + assert_eq!(selected.len(), 2); + assert_eq!( + selected + .iter() + .filter_map(|window| window.source_id.as_ref()) + .collect::>() + .len(), + 2 + ); + } + + #[test] + fn raw_quota_history_isolated_to_unattributed_window_source() { + let window = test_unattributed_quota_window("source-a", "window-a"); + let observations = raw_observations_for_window( + vec![ + test_unattributed_quota_record("source-a"), + test_unattributed_quota_record("source-b"), + ], + &window, + ); + + assert_eq!(observations.len(), 1); + assert_eq!(observations[0].observation.source_id.0, "source-a"); + } } diff --git a/crates/statsai/tests/quota_cli.rs b/crates/statsai/tests/quota_cli.rs new file mode 100644 index 0000000..f0660ba --- /dev/null +++ b/crates/statsai/tests/quota_cli.rs @@ -0,0 +1,432 @@ +use serde_json::Value; +use std::fs::File; +use std::io::Write; +use std::path::Path; +use std::process::{Command, Output}; +use tempfile::tempdir; + +fn run_statsai(store: &Path, codex_home: &Path, args: &[&str]) -> Output { + Command::new(env!("CARGO_BIN_EXE_statsai")) + .arg("--store") + .arg(store) + .arg("--device-id") + .arg("quota-device") + .args(args) + .env("CODEX_HOME", codex_home) + .output() + .expect("run statsai") +} + +fn stdout(output: Output) -> String { + assert!( + output.status.success(), + "command failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + String::from_utf8(output.stdout).expect("utf-8 stdout") +} + +#[test] +fn quota_cli_scans_attributes_reconstructs_and_exports_sync_projection() { + let directory = tempdir().expect("temp dir"); + let codex_home = directory.path().join("codex"); + let sessions = codex_home.join("sessions/2026/08/20"); + std::fs::create_dir_all(&sessions).expect("sessions"); + let mut fixture = File::create(sessions.join("thread.jsonl")).expect("fixture"); + writeln!( + fixture, + "{}", + serde_json::json!({ + "timestamp": "2026-08-20T12:00:00Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15 + }, + "total_token_usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15 + } + }, + "rate_limits": { + "limit_id": "weekly,\"special\"", + "plan_type": "pro", + "primary": { + "used_percent": 25, + "window_minutes": 10080, + "resets_at": 1787832000 + }, + "secondary": { + "used_percent": 5, + "window_minutes": 300, + "resets_at": 1787245200 + }, + "credits": {"balance": "10.50"} + } + } + }) + ) + .expect("write fixture"); + writeln!( + fixture, + "{}", + serde_json::json!({ + "timestamp": "2026-08-20T12:01:00Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "rate_limits": { + "limit_id": "weekly,\"special\"", + "plan_type": "pro", + "primary": { + "used_percent": 6, + "window_minutes": 300, + "resets_at": 1787245320 + } + } + } + }) + ) + .expect("write current single-window fixture"); + drop(fixture); + let store = directory.path().join("statsai.sqlite3"); + + let scan = stdout(run_statsai( + &store, + &codex_home, + &["scan", "--provider", "codex"], + )); + assert!(scan.contains("quota_observations=2"), "{scan}"); + + let status: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "status", "--json"], + ))) + .expect("status JSON"); + assert_eq!(status["total_observations"], 2); + assert_eq!(status["unattributed_observations"], 2); + + let history: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "history", "--raw", "--json"], + ))) + .expect("raw history JSON"); + assert_eq!( + history["observations"][0]["observation"]["usage_link_kind"], + "record_event" + ); + assert!(history["observations"][0]["observation"]["usage_event_id"].is_string()); + + let windows: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "windows", "--json"], + ))) + .expect("windows JSON"); + assert_eq!(windows.as_array().expect("windows").len(), 2); + assert!(windows + .as_array() + .expect("windows") + .iter() + .all(|window| window["usage_totals"].is_null())); + let filtered_windows: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &[ + "quota", + "windows", + "--from", + "2026-08-20T12:01:00Z", + "--json", + ], + ))) + .expect("filtered windows JSON"); + assert_eq!( + filtered_windows.as_array().expect("filtered windows").len(), + 1 + ); + let filtered_window_id = filtered_windows[0]["window_id"] + .as_str() + .expect("filtered window id"); + let full_window_id = windows + .as_array() + .expect("windows") + .iter() + .find(|window| window["window_minutes"] == 300) + .and_then(|window| window["window_id"].as_str()) + .expect("full window id"); + assert_eq!(filtered_window_id, full_window_id); + let filtered_history: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &[ + "quota", + "history", + "--window-id", + filtered_window_id, + "--json", + ], + ))) + .expect("filtered window history JSON"); + assert_eq!(filtered_history["window_id"], filtered_window_id); + let human_windows = stdout(run_statsai(&store, &codex_home, &["quota", "windows"])); + assert!( + human_windows.contains("usage unavailable"), + "{human_windows}" + ); + + let no_windows: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "windows", "--limit", "0", "--json"], + ))) + .expect("zero-limit windows JSON"); + assert_eq!(no_windows, serde_json::json!([])); + + let current: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "current", "--json"], + ))) + .expect("current JSON"); + assert_eq!(current.as_array().expect("current").len(), 1); + assert_eq!(current[0]["window_minutes"], 10_080); + + stdout(run_statsai( + &store, + &codex_home, + &[ + "source", + "connect", + "--path", + codex_home.to_str().expect("path"), + "--provider-user-id", + "account-cli", + "--started-at", + "2026-08-01", + ], + )); + + let attributed_windows: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "windows", "--json"], + ))) + .expect("attributed windows JSON"); + let weekly_window = attributed_windows + .as_array() + .expect("attributed windows") + .iter() + .find(|window| window["window_minutes"] == 10_080) + .expect("weekly window"); + assert_eq!(weekly_window["usage_totals"]["total_tokens"], 15); + + let projections: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &[ + "quota", + "export", + "--level", + "sync-windows", + "--format", + "json", + ], + ))) + .expect("projection JSON"); + assert_eq!(projections.as_array().expect("projections").len(), 1); + assert_eq!(projections[0]["device_id"], "quota-device"); + assert!(projections[0].get("source_id").is_none()); + assert!(projections[0].get("total_tokens").is_none()); + + let csv = stdout(run_statsai( + &store, + &codex_home, + &["quota", "export", "--level", "windows", "--format", "csv"], + )); + let mut reader = csv::Reader::from_reader(csv.as_bytes()); + let rows = reader + .records() + .collect::>>() + .expect("standards-compliant CSV"); + assert_eq!(rows.len(), 2); + assert!(rows + .iter() + .all(|row| row.get(5) == Some("weekly,\"special\""))); +} + +#[test] +fn quota_projection_schema_is_available_without_opening_a_store() { + let directory = tempdir().expect("temp dir"); + let missing_store = directory.path().join("missing/store.sqlite3"); + let schema: Value = serde_json::from_str(&stdout(run_statsai( + &missing_store, + directory.path(), + &["schema", "quota-window-projection"], + ))) + .expect("schema JSON"); + assert_eq!(schema["properties"]["schema_version"]["type"], "string"); + assert!(!missing_store.exists()); +} + +#[test] +fn scan_preview_reports_quota_only_sources_and_totals() { + let directory = tempdir().expect("temp dir"); + let codex_home = directory.path().join("codex"); + let sessions = codex_home.join("sessions/2026/08/20"); + std::fs::create_dir_all(&sessions).expect("sessions"); + let mut fixture = File::create(sessions.join("quota-only.jsonl")).expect("fixture"); + writeln!( + fixture, + "{}", + serde_json::json!({ + "timestamp": "2026-08-20T12:00:00Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "rate_limits": { + "primary": { + "used_percent": 25, + "window_minutes": 10080, + "resets_at": 1787832000 + } + } + } + }) + ) + .expect("write fixture"); + + let preview = stdout(run_statsai( + &directory.path().join("statsai.sqlite3"), + &codex_home, + &["scan", "--provider", "codex", "--preview"], + )); + + assert!( + preview.contains("usage_events=0 summaries=0 task_spans=0 quota_observations=1"), + "{preview}" + ); + assert!( + preview + .contains("preview total: sources=1 usage_events=0 summaries=0 quota_observations=1"), + "{preview}" + ); +} + +#[test] +fn conversation_archive_backfills_and_reconciles_quota_observations() { + let directory = tempdir().expect("temp dir"); + let codex_home = directory.path().join("codex"); + let sessions = codex_home.join("sessions"); + std::fs::create_dir_all(&sessions).expect("sessions"); + let path = sessions.join("archived-thread.jsonl"); + let mut fixture = File::create(&path).expect("fixture"); + writeln!( + fixture, + "{}", + serde_json::json!({ + "timestamp": "2026-08-20T12:00:00Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "rate_limits": { + "primary": { + "used_percent": 25, + "window_minutes": 10080, + "resets_at": 1787832000 + } + } + } + }) + ) + .expect("write fixture"); + writeln!( + fixture, + "{}", + serde_json::json!({ + "timestamp": "2026-08-20T12:05:00Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "rate_limits": { + "primary": { + "used_percent": 30, + "window_minutes": 10080, + "resets_at": 1787832000 + } + } + } + }) + ) + .expect("write second fixture record"); + drop(fixture); + let store = directory.path().join("statsai.sqlite3"); + + stdout(run_statsai( + &store, + &codex_home, + &["conversation", "collect", "--provider", "codex"], + )); + let status: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "status", "--json"], + ))) + .expect("status JSON"); + assert_eq!(status["total_observations"], 2); + + let mut fixture = File::create(&path).expect("replace fixture"); + writeln!( + fixture, + "{}", + serde_json::json!({ + "timestamp": "2026-08-20T12:00:00Z", + "type": "event_msg", + "payload": { + "type": "token_count", + "rate_limits": { + "primary": { + "used_percent": 25, + "window_minutes": 10080, + "resets_at": 1787832000 + } + } + } + }) + ) + .expect("write shortened fixture"); + drop(fixture); + stdout(run_statsai( + &store, + &codex_home, + &["conversation", "collect", "--provider", "codex"], + )); + let status: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "status", "--json"], + ))) + .expect("shortened status JSON"); + assert_eq!(status["total_observations"], 1); + + std::fs::remove_file(path).expect("remove archive fixture"); + stdout(run_statsai( + &store, + &codex_home, + &["conversation", "collect", "--provider", "codex"], + )); + let status: Value = serde_json::from_str(&stdout(run_statsai( + &store, + &codex_home, + &["quota", "status", "--json"], + ))) + .expect("reconciled status JSON"); + assert_eq!(status["total_observations"], 0); +} diff --git a/docs/fixtures/quota_window_projection.v1.two_devices.json b/docs/fixtures/quota_window_projection.v1.two_devices.json new file mode 100644 index 0000000..86205bf --- /dev/null +++ b/docs/fixtures/quota_window_projection.v1.two_devices.json @@ -0,0 +1,90 @@ +[ + { + "schema_version": "quota_window_sync_projection.v1", + "projection_id": "quota_projection_device_a", + "device_id": "device-a", + "provider": "codex", + "provider_account_id": "acct_stable", + "limit_id": "subscription", + "window_minutes": 10080, + "inferred_start": "2026-08-17T12:00:00Z", + "representative_reset": "2026-08-24T12:00:00Z", + "representative_reset_epoch_seconds": 1787572800, + "reset_min": "2026-08-24T12:00:00Z", + "reset_min_epoch_seconds": 1787572800, + "reset_max": "2026-08-24T12:02:00Z", + "reset_max_epoch_seconds": 1787572920, + "first_observed_at": "2026-08-20T09:00:00Z", + "last_observed_at": "2026-08-20T10:00:00Z", + "sample_count": 2, + "first_used_percent": 10.0, + "latest_used_percent": 15.0, + "minimum_used_percent": 10.0, + "maximum_used_percent": 15.0, + "change_points": [ + { + "observed_at": "2026-08-20T09:00:00Z", + "used_percent": 10.0, + "resets_at": "2026-08-24T12:00:00Z", + "resets_at_epoch_seconds": 1787572800, + "point_fingerprint": "point_shared_10", + "provider_slot": "primary" + } + ], + "latest_status": { + "plan_type": "pro", + "individual_limit": null, + "spend_control_state": null, + "reached_type": null, + "credits": { + "has_credits": true, + "unlimited": false, + "balance": "10.5" + } + } + }, + { + "schema_version": "quota_window_sync_projection.v1", + "projection_id": "quota_projection_device_b", + "device_id": "device-b", + "provider": "codex", + "provider_account_id": "acct_stable", + "limit_id": "subscription", + "window_minutes": 10080, + "inferred_start": "2026-08-17T12:03:00Z", + "representative_reset": "2026-08-24T12:03:00Z", + "representative_reset_epoch_seconds": 1787572980, + "reset_min": "2026-08-24T12:03:00Z", + "reset_min_epoch_seconds": 1787572980, + "reset_max": "2026-08-24T12:03:00Z", + "reset_max_epoch_seconds": 1787572980, + "first_observed_at": "2026-08-20T09:00:00Z", + "last_observed_at": "2026-08-20T11:00:00Z", + "sample_count": 2, + "first_used_percent": 10.0, + "latest_used_percent": 18.0, + "minimum_used_percent": 10.0, + "maximum_used_percent": 18.0, + "change_points": [ + { + "observed_at": "2026-08-20T09:00:00Z", + "used_percent": 10.0, + "resets_at": "2026-08-24T12:00:00Z", + "resets_at_epoch_seconds": 1787572800, + "point_fingerprint": "point_shared_10", + "provider_slot": "secondary" + } + ], + "latest_status": { + "plan_type": "pro", + "individual_limit": null, + "spend_control_state": null, + "reached_type": null, + "credits": { + "has_credits": true, + "unlimited": false, + "balance": "10.5" + } + } + } +] diff --git a/docs/quota-sync.md b/docs/quota-sync.md new file mode 100644 index 0000000..8eaedc3 --- /dev/null +++ b/docs/quota-sync.md @@ -0,0 +1,41 @@ +# Quota window projection contract + +StatsAI stores provider quota observations locally and reconstructs reset windows without adding +them to `SyncBatch.v3`. The standalone hosted-facing contract is available from the CLI: + +```sh +statsai schema quota-window-projection +``` + +`quota_window_sync_projection.v1` is a device contribution. Its `projection_id` is deterministic +for the device, provider account, limit identity, duration, and cluster anchor point. It is not a +logical account-window identifier. Projections omit source paths, local source/account row +identities, provider payloads, usage tokens, and cost. + +Only attributed windows lasting at least 10,080 minutes are emitted by: + +```sh +statsai quota export --level sync-windows --format jsonl +``` + +The command warns when locally observed weekly history is not fully attributed. + +Before attribution, local window reconstruction is partitioned by stable source ID so evidence +from separate local provider installations cannot be blended. That source partition is local-only: +once evidence is attributed, account identity becomes the reconstruction scope, and no source ID is +included in the sync projection. + +## Server merge semantics + +1. Group contributions by stable provider account, provider, limit identity, and duration. +2. Cluster reset epochs whose raw reset values are within five minutes. This tolerance covers + provider recomputation drift; it is not a device clock-skew correction. +3. Deduplicate change points by `point_fingerprint`. +4. Preserve the chronological percentage sequence. Do not sum percentages and do not replace the + sequence with a maximum observed percentage. +5. Retain the logical window while at least one device projection contributes it. +6. Join the merged interval to the existing server-deduplicated usage rollup for the provider + account. Never derive or sum token/cost values from quota projections. + +Adding this projection to a future hosted wire payload is a separate capability and versioning +decision; this contract does not change `SyncBatch.v3`.