1use crate::{DbError, DbResult};
4use projectatlas_core::graph::ProjectInstanceId;
5use projectatlas_core::telemetry::{
6 TOKEN_AVERAGE_POLICY_OVERFLOW_EVIDENCE, TOKEN_BASELINE_DIRECTORY_WALK, TokenAccountingTotals,
7 TokenBucketOverview, TokenOverview, TokenTrendPeriod, TokenTrendReport, TokenTrendWindow,
8 UsageDetailAvailability, UsageEvent, UsageInstanceId, UsageInstanceOwner,
9 average_modeled_baseline_tokens,
10};
11use rusqlite::{Connection, OptionalExtension, params};
12use serde::Serialize;
13use std::collections::{BTreeMap, BTreeSet};
14use std::time::{Duration, SystemTime, UNIX_EPOCH};
15
16const POLICY_VERSION: u32 = 1;
18const LOGICAL_BYTE_VERSION: u32 = 1;
20const OVERFLOW_DIMENSION: &str = "<overflow>";
22const INSTANCE_ACTIVE: &str = "active";
24const INSTANCE_SEALED: &str = "sealed";
26const INSTANCE_EXPIRED: &str = "expired";
28const DEDUPE_SCOPE_EVENT: &str = "event";
30const SECONDS_PER_DAY: i64 = 86_400;
32const AGGREGATE_COUNTER_FIELD: &str = "aggregate_counter";
34const LEGACY_TEXT_HASH_DOMAIN: &[u8] = b"projectatlas:legacy-telemetry-text:v1\0";
36const WORKTREE_USAGE_ROUTED: &str = "routed";
38const WORKTREE_USAGE_SYNCHRONIZED: &str = "synchronized";
40const MAX_WORKTREE_USAGE_SNAPSHOT_BYTES: usize = 32 * 1_024 * 1_024;
42
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
45enum BaselineAdmission {
46 BoundedRuntime,
48 SupportedUpgrade,
50}
51
52#[derive(Clone, Copy, Debug, Eq, PartialEq)]
54enum DimensionAdmission {
55 Event,
57 Overflow,
59}
60
61#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
63struct LegacyDetailLoss {
64 raw: bool,
66 dimension: bool,
68 label: bool,
70}
71
72macro_rules! aggregate_params {
74 ($first:expr, $second:expr, $value:expr) => {
75 params![
76 $first,
77 $second,
78 $value.calls,
79 $value.estimated_without,
80 $value.estimated_with,
81 $value.observed_without,
82 $value.observed_with,
83 $value.modeled_without,
84 $value.modeled_with,
85 $value.deduped_modeled_without,
86 $value.deduped_modeled_with,
87 $value.repeated_baselines,
88 $value.observed_file_read_replacements,
89 $value.modeled_file_reads_avoided,
90 ]
91 };
92}
93
94macro_rules! daily_aggregate_params {
96 ($first:expr, $day:expr, $dimension:expr, $value:expr) => {
97 params![
98 $first,
99 $day,
100 $dimension,
101 $value.calls,
102 $value.estimated_without,
103 $value.estimated_with,
104 $value.observed_without,
105 $value.observed_with,
106 $value.modeled_without,
107 $value.modeled_with,
108 $value.deduped_modeled_without,
109 $value.deduped_modeled_with,
110 $value.repeated_baselines,
111 $value.observed_file_read_replacements,
112 $value.modeled_file_reads_avoided,
113 ]
114 };
115}
116
117macro_rules! worktree_aggregate_params {
119 ($registration:expr, $source:expr, $day:expr, $dimension:expr, $value:expr) => {
120 params![
121 $registration,
122 $source,
123 $day,
124 $dimension,
125 $value.calls,
126 $value.estimated_without,
127 $value.estimated_with,
128 $value.observed_without,
129 $value.observed_with,
130 $value.modeled_without,
131 $value.modeled_with,
132 $value.deduped_modeled_without,
133 $value.deduped_modeled_with,
134 $value.repeated_baselines,
135 $value.observed_file_read_replacements,
136 $value.modeled_file_reads_avoided,
137 ]
138 };
139}
140
141pub(crate) fn generate_usage_instance_id() -> DbResult<UsageInstanceId> {
143 let mut bytes = [0_u8; 16];
144 getrandom::fill(&mut bytes).map_err(|_source| DbError::TelemetryIdentityUnavailable)?;
145 UsageInstanceId::from_bytes(bytes).map_err(Into::into)
146}
147
148#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
150pub struct TelemetryRetentionPolicy {
151 pub max_raw_rows: usize,
153 pub max_raw_logical_bytes: usize,
155 pub max_raw_age_seconds: u64,
157 pub max_dimensions: usize,
162 pub max_active_instances: usize,
164 pub max_retained_instances: usize,
166 pub max_label_tombstones: usize,
168 pub max_instance_tombstones: usize,
170 pub max_retained_labels: usize,
172 pub max_baselines_per_instance: usize,
174 pub max_active_baseline_rows: usize,
176 pub max_baseline_logical_bytes: usize,
178 pub max_daily_rows: usize,
180 pub prune_batch_rows: usize,
182 pub max_label_bytes: usize,
184 pub max_command_bytes: usize,
186 pub max_path_bytes: usize,
188 pub max_query_bytes: usize,
190 pub max_dimension_bytes: usize,
192 pub max_baseline_witness_bytes: usize,
194 pub max_active_idle_seconds: u64,
196 pub future_clock_tolerance_seconds: u64,
198 pub checkpoint_write_interval: usize,
200 pub retained_trend_days: u64,
202 pub retained_instance_seconds: u64,
204 pub retained_label_seconds: u64,
206 pub retained_tombstone_seconds: u64,
208}
209
210impl Default for TelemetryRetentionPolicy {
211 fn default() -> Self {
212 Self {
213 max_raw_rows: 50_000,
214 max_raw_logical_bytes: 64 * 1_024 * 1_024,
215 max_raw_age_seconds: 30 * 24 * 60 * 60,
216 max_dimensions: 128,
217 max_active_instances: 64,
218 max_retained_instances: 4_096,
219 max_label_tombstones: 1_024,
220 max_instance_tombstones: 4_096,
221 max_retained_labels: 256,
222 max_baselines_per_instance: 1_024,
223 max_active_baseline_rows: 16_384,
224 max_baseline_logical_bytes: 16 * 1_024 * 1_024,
225 max_daily_rows: 100_000,
226 prune_batch_rows: 512,
227 max_label_bytes: 128,
228 max_command_bytes: 96,
229 max_path_bytes: 4_096,
230 max_query_bytes: 4_096,
231 max_dimension_bytes: 96,
232 max_baseline_witness_bytes: 1_024,
233 max_active_idle_seconds: 24 * 60 * 60,
234 future_clock_tolerance_seconds: 5 * 60,
235 checkpoint_write_interval: 1_024,
236 retained_trend_days: 400,
237 retained_instance_seconds: 90 * 24 * 60 * 60,
238 retained_label_seconds: 365 * 24 * 60 * 60,
239 retained_tombstone_seconds: 730 * 24 * 60 * 60,
240 }
241 }
242}
243
244impl TelemetryRetentionPolicy {
245 pub fn validate(self) -> DbResult<Self> {
251 for (field, value) in [
252 ("max_raw_rows", self.max_raw_rows),
253 ("max_raw_logical_bytes", self.max_raw_logical_bytes),
254 ("max_dimensions", self.max_dimensions),
255 ("max_active_instances", self.max_active_instances),
256 ("max_retained_instances", self.max_retained_instances),
257 ("max_label_tombstones", self.max_label_tombstones),
258 ("max_instance_tombstones", self.max_instance_tombstones),
259 ("max_retained_labels", self.max_retained_labels),
260 (
261 "max_baselines_per_instance",
262 self.max_baselines_per_instance,
263 ),
264 ("max_active_baseline_rows", self.max_active_baseline_rows),
265 (
266 "max_baseline_logical_bytes",
267 self.max_baseline_logical_bytes,
268 ),
269 ("max_daily_rows", self.max_daily_rows),
270 ("prune_batch_rows", self.prune_batch_rows),
271 ("max_label_bytes", self.max_label_bytes),
272 ("max_command_bytes", self.max_command_bytes),
273 ("max_path_bytes", self.max_path_bytes),
274 ("max_query_bytes", self.max_query_bytes),
275 ("max_dimension_bytes", self.max_dimension_bytes),
276 (
277 "max_baseline_witness_bytes",
278 self.max_baseline_witness_bytes,
279 ),
280 ("checkpoint_write_interval", self.checkpoint_write_interval),
281 ] {
282 if value == 0 {
283 return Err(DbError::TelemetryLimitInvalid { field, value });
284 }
285 }
286 if self.max_raw_age_seconds == 0
287 || self.max_active_idle_seconds == 0
288 || self.retained_trend_days == 0
289 || self.retained_instance_seconds == 0
290 || self.retained_label_seconds == 0
291 || self.retained_tombstone_seconds == 0
292 || self.max_retained_instances < self.max_active_instances
293 || self.max_daily_rows < 2
294 {
295 return Err(DbError::TelemetryLimitInvalid {
296 field: "telemetry_retention_policy",
297 value: 0,
298 });
299 }
300 Ok(self)
301 }
302}
303
304#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
306#[serde(rename_all = "snake_case")]
307pub enum SpillCleanupState {
308 NotApplicable,
310}
311
312#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
314#[serde(rename_all = "snake_case")]
315pub enum TelemetryCheckpointState {
316 NotDue,
318 Completed,
320 Busy,
322 Error,
324}
325
326impl TelemetryCheckpointState {
327 const fn as_str(self) -> &'static str {
329 match self {
330 Self::NotDue => "not_due",
331 Self::Completed => "completed",
332 Self::Busy => "busy",
333 Self::Error => "error",
334 }
335 }
336
337 fn from_str(value: &str) -> DbResult<Self> {
339 match value {
340 "not_due" => Ok(Self::NotDue),
341 "completed" => Ok(Self::Completed),
342 "busy" => Ok(Self::Busy),
343 "error" => Ok(Self::Error),
344 _ => Err(DbError::InvalidEnum {
345 field: "usage_retention_state.checkpoint_state",
346 value: value.to_string(),
347 }),
348 }
349 }
350}
351
352#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
354#[serde(rename_all = "snake_case")]
355pub enum PlannerStatisticsPolicy {
356 NotConfigured,
358}
359
360#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
362#[serde(rename_all = "snake_case")]
363pub enum PlannerStatisticsState {
364 NotInitialized,
366 Available,
368}
369
370#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
372pub struct TelemetryRetentionState {
373 pub policy_version: u32,
375 pub logical_byte_version: u32,
377 pub raw_rows: usize,
379 pub max_raw_rows: usize,
381 pub max_raw_age_seconds: u64,
383 pub raw_logical_bytes: usize,
385 pub max_raw_logical_bytes: usize,
387 pub baseline_rows: usize,
389 pub max_baselines_per_instance: usize,
391 pub max_active_baseline_rows: usize,
393 pub baseline_logical_bytes: usize,
395 pub max_baseline_logical_bytes: usize,
397 pub dimension_rows: usize,
399 pub max_dimensions: usize,
401 pub instance_rows: usize,
403 pub active_instance_rows: usize,
405 pub max_active_instances: usize,
407 pub max_retained_instances: usize,
409 pub retained_label_rows: usize,
411 pub max_retained_labels: usize,
413 pub daily_rows: usize,
415 pub max_daily_rows: usize,
417 pub retained_trend_days: u64,
419 pub label_tombstone_rows: usize,
421 pub max_label_tombstones: usize,
423 pub instance_tombstone_rows: usize,
425 pub max_instance_tombstones: usize,
427 pub pruned_raw_rows: usize,
429 pub pruned_instance_rows: usize,
431 pub evicted_tombstones: usize,
433 pub maintenance_pending: bool,
435 pub prune_batch_rows: usize,
437 pub writes_since_checkpoint: usize,
439 pub checkpoint_write_interval: usize,
441 pub last_checkpoint_epoch: u64,
443 pub oldest_retained_epoch: Option<u64>,
445 pub clock_anomaly: bool,
447 pub spill_cleanup: SpillCleanupState,
449 pub checkpoint_state: TelemetryCheckpointState,
451 pub wal_autocheckpoint_pages: usize,
453 pub freelist_pages: usize,
455 pub page_count: usize,
457 pub page_size: usize,
459 pub journal_mode: String,
461 pub synchronous_mode: String,
463 pub connection_busy_timeout_ms: u64,
465 pub normal_busy_timeout_ms: u64,
467 pub telemetry_busy_timeout_ms: u64,
469 pub statistics_policy: PlannerStatisticsPolicy,
471 pub statistics_state: PlannerStatisticsState,
473}
474
475#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
477struct DimensionValues {
478 token_savings_bucket: String,
480 provider: String,
482 model: String,
484 tokenizer_backend: String,
486 accuracy: String,
488 baseline_kind: String,
490 confidence: String,
492 accounting_layer: String,
494 estimate_method: String,
496 denominator_kind: String,
498 dedupe_scope: String,
500 overflow: bool,
502}
503
504impl DimensionValues {
505 fn from_event(event: &UsageEvent) -> Self {
507 Self {
508 token_savings_bucket: event.token_savings_bucket.clone(),
509 provider: event.provider.clone(),
510 model: event.model.clone(),
511 tokenizer_backend: event.tokenizer_backend.clone(),
512 accuracy: event.accuracy.clone(),
513 baseline_kind: event.baseline_kind.clone(),
514 confidence: event.confidence.clone(),
515 accounting_layer: event.report_accounting_layer().to_string(),
516 estimate_method: event.estimate_method.clone(),
517 denominator_kind: event.report_denominator_kind().to_string(),
518 dedupe_scope: event.report_dedupe_scope().to_string(),
519 overflow: false,
520 }
521 }
522
523 fn overflow(denominator_kind: &str) -> Self {
525 Self {
526 token_savings_bucket: OVERFLOW_DIMENSION.to_string(),
527 provider: OVERFLOW_DIMENSION.to_string(),
528 model: OVERFLOW_DIMENSION.to_string(),
529 tokenizer_backend: OVERFLOW_DIMENSION.to_string(),
530 accuracy: OVERFLOW_DIMENSION.to_string(),
531 baseline_kind: OVERFLOW_DIMENSION.to_string(),
532 confidence: OVERFLOW_DIMENSION.to_string(),
533 accounting_layer: OVERFLOW_DIMENSION.to_string(),
534 estimate_method: OVERFLOW_DIMENSION.to_string(),
535 denominator_kind: if denominator_kind == TOKEN_BASELINE_DIRECTORY_WALK {
536 TOKEN_BASELINE_DIRECTORY_WALK.to_string()
537 } else {
538 OVERFLOW_DIMENSION.to_string()
539 },
540 dedupe_scope: OVERFLOW_DIMENSION.to_string(),
541 overflow: true,
542 }
543 }
544}
545
546#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
548struct AggregateCounters {
549 calls: i64,
551 estimated_without: i64,
553 estimated_with: i64,
555 observed_without: i64,
557 observed_with: i64,
559 modeled_without: i64,
561 modeled_with: i64,
563 deduped_modeled_without: i64,
565 deduped_modeled_with: i64,
567 repeated_baselines: i64,
569 observed_file_read_replacements: i64,
571 modeled_file_reads_avoided: i64,
573}
574
575#[derive(Clone, Debug, Eq, PartialEq)]
577struct WorktreeUsageSnapshotRow {
578 day_epoch: i64,
580 dimension_id: i64,
582 counters: AggregateCounters,
584}
585
586#[derive(Clone, Debug, Eq, PartialEq)]
588pub struct WorktreeUsageSnapshot {
589 project_instance_id: ProjectInstanceId,
591 revision: u64,
593 dimensions: BTreeMap<i64, DimensionValues>,
595 rows: Vec<WorktreeUsageSnapshotRow>,
597 logical_bytes: usize,
599}
600
601impl WorktreeUsageSnapshot {
602 #[must_use]
604 pub const fn project_instance_id(&self) -> ProjectInstanceId {
605 self.project_instance_id
606 }
607
608 #[must_use]
610 pub const fn revision(&self) -> u64 {
611 self.revision
612 }
613
614 #[must_use]
616 pub fn row_count(&self) -> usize {
617 self.rows.len()
618 }
619
620 #[must_use]
622 pub const fn logical_bytes(&self) -> usize {
623 self.logical_bytes
624 }
625}
626
627#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
629#[serde(rename_all = "snake_case")]
630pub enum WorktreeUsageSyncState {
631 Synchronized,
633 Current,
635}
636
637#[derive(Clone, Copy, Debug)]
639enum RetentionCounter {
640 RawRows,
642 RawLogicalBytes,
644 BaselineRows,
646 BaselineLogicalBytes,
648 DimensionRows,
650 InstanceRows,
652 LabelRows,
654 DailyRows,
656 LabelTombstoneRows,
658 InstanceTombstoneRows,
660}
661
662impl RetentionCounter {
663 const fn select_sql(self) -> &'static str {
665 match self {
666 Self::RawRows => "SELECT raw_rows FROM usage_retention_state WHERE singleton = 1",
667 Self::RawLogicalBytes => {
668 "SELECT raw_logical_bytes FROM usage_retention_state WHERE singleton = 1"
669 }
670 Self::BaselineRows => {
671 "SELECT baseline_rows FROM usage_retention_state WHERE singleton = 1"
672 }
673 Self::BaselineLogicalBytes => {
674 "SELECT baseline_logical_bytes FROM usage_retention_state WHERE singleton = 1"
675 }
676 Self::DimensionRows => {
677 "SELECT dimension_rows FROM usage_retention_state WHERE singleton = 1"
678 }
679 Self::InstanceRows => {
680 "SELECT instance_rows FROM usage_retention_state WHERE singleton = 1"
681 }
682 Self::LabelRows => "SELECT label_rows FROM usage_retention_state WHERE singleton = 1",
683 Self::DailyRows => "SELECT daily_rows FROM usage_retention_state WHERE singleton = 1",
684 Self::LabelTombstoneRows => {
685 "SELECT label_tombstone_rows FROM usage_retention_state WHERE singleton = 1"
686 }
687 Self::InstanceTombstoneRows => {
688 "SELECT instance_tombstone_rows FROM usage_retention_state WHERE singleton = 1"
689 }
690 }
691 }
692
693 const fn update_sql(self) -> &'static str {
695 match self {
696 Self::RawRows => "UPDATE usage_retention_state SET raw_rows = ?1 WHERE singleton = 1",
697 Self::RawLogicalBytes => {
698 "UPDATE usage_retention_state SET raw_logical_bytes = ?1 WHERE singleton = 1"
699 }
700 Self::BaselineRows => {
701 "UPDATE usage_retention_state SET baseline_rows = ?1 WHERE singleton = 1"
702 }
703 Self::BaselineLogicalBytes => {
704 "UPDATE usage_retention_state SET baseline_logical_bytes = ?1 WHERE singleton = 1"
705 }
706 Self::DimensionRows => {
707 "UPDATE usage_retention_state SET dimension_rows = ?1 WHERE singleton = 1"
708 }
709 Self::InstanceRows => {
710 "UPDATE usage_retention_state SET instance_rows = ?1 WHERE singleton = 1"
711 }
712 Self::LabelRows => {
713 "UPDATE usage_retention_state SET label_rows = ?1 WHERE singleton = 1"
714 }
715 Self::DailyRows => {
716 "UPDATE usage_retention_state SET daily_rows = ?1 WHERE singleton = 1"
717 }
718 Self::LabelTombstoneRows => {
719 "UPDATE usage_retention_state SET label_tombstone_rows = ?1 WHERE singleton = 1"
720 }
721 Self::InstanceTombstoneRows => {
722 "UPDATE usage_retention_state SET instance_tombstone_rows = ?1 WHERE singleton = 1"
723 }
724 }
725 }
726
727 const fn field(self) -> &'static str {
729 match self {
730 Self::RawRows => "raw_rows",
731 Self::RawLogicalBytes => "raw_logical_bytes",
732 Self::BaselineRows => "baseline_rows",
733 Self::BaselineLogicalBytes => "baseline_logical_bytes",
734 Self::DimensionRows => "dimension_rows",
735 Self::InstanceRows => "instance_rows",
736 Self::LabelRows => "label_rows",
737 Self::DailyRows => "daily_rows",
738 Self::LabelTombstoneRows => "label_tombstone_rows",
739 Self::InstanceTombstoneRows => "instance_tombstone_rows",
740 }
741 }
742}
743
744impl AggregateCounters {
745 fn validate_nonnegative(self) -> DbResult<()> {
747 for (field, value) in [
748 ("calls", self.calls),
749 ("estimated_without", self.estimated_without),
750 ("estimated_with", self.estimated_with),
751 ("observed_without", self.observed_without),
752 ("observed_with", self.observed_with),
753 ("modeled_without", self.modeled_without),
754 ("modeled_with", self.modeled_with),
755 ("deduped_modeled_without", self.deduped_modeled_without),
756 ("deduped_modeled_with", self.deduped_modeled_with),
757 ("repeated_baselines", self.repeated_baselines),
758 (
759 "observed_file_read_replacements",
760 self.observed_file_read_replacements,
761 ),
762 (
763 "modeled_file_reads_avoided",
764 self.modeled_file_reads_avoided,
765 ),
766 ] {
767 if value < 0 {
768 return Err(DbError::TelemetryIntegerOverflow { field });
769 }
770 }
771 Ok(())
772 }
773
774 fn checked_add(self, other: Self) -> DbResult<Self> {
776 macro_rules! add {
777 ($field:ident) => {
778 self.$field
779 .checked_add(other.$field)
780 .ok_or(DbError::TelemetryIntegerOverflow {
781 field: stringify!($field),
782 })?
783 };
784 }
785 Ok(Self {
786 calls: add!(calls),
787 estimated_without: add!(estimated_without),
788 estimated_with: add!(estimated_with),
789 observed_without: add!(observed_without),
790 observed_with: add!(observed_with),
791 modeled_without: add!(modeled_without),
792 modeled_with: add!(modeled_with),
793 deduped_modeled_without: add!(deduped_modeled_without),
794 deduped_modeled_with: add!(deduped_modeled_with),
795 repeated_baselines: add!(repeated_baselines),
796 observed_file_read_replacements: add!(observed_file_read_replacements),
797 modeled_file_reads_avoided: add!(modeled_file_reads_avoided),
798 })
799 }
800
801 fn contains(self, accepted: Self) -> bool {
803 [
804 (self.calls, accepted.calls),
805 (self.estimated_without, accepted.estimated_without),
806 (self.estimated_with, accepted.estimated_with),
807 (self.observed_without, accepted.observed_without),
808 (self.observed_with, accepted.observed_with),
809 (self.modeled_without, accepted.modeled_without),
810 (self.modeled_with, accepted.modeled_with),
811 (
812 self.deduped_modeled_without,
813 accepted.deduped_modeled_without,
814 ),
815 (self.deduped_modeled_with, accepted.deduped_modeled_with),
816 (self.repeated_baselines, accepted.repeated_baselines),
817 (
818 self.observed_file_read_replacements,
819 accepted.observed_file_read_replacements,
820 ),
821 (
822 self.modeled_file_reads_avoided,
823 accepted.modeled_file_reads_avoided,
824 ),
825 ]
826 .into_iter()
827 .all(|(incoming, accepted)| incoming >= accepted)
828 }
829}
830
831pub(crate) fn initialize_empty_storage(connection: &Connection) -> DbResult<()> {
833 let policy = TelemetryRetentionPolicy::default().validate()?;
834 ensure_overflow_dimension(connection, OVERFLOW_DIMENSION)?;
835 refresh_retention_state(connection, policy, now_epoch_seconds()?, 0, 0, 0, 0)
836}
837
838pub(crate) fn reset_usage_storage_for_hydration(connection: &Connection) -> DbResult<()> {
840 connection.execute_batch(
841 "DELETE FROM usage_events;
842 DELETE FROM usage_instances;
843 DELETE FROM usage_global_aggregates;
844 DELETE FROM usage_daily_aggregates;
845 DELETE FROM usage_labels;
846 DELETE FROM usage_label_tombstones;
847 DELETE FROM usage_instance_tombstones;
848 DELETE FROM usage_bucket_dimensions;
849 UPDATE usage_retention_state SET
850 raw_rows = 0,
851 raw_logical_bytes = 0,
852 baseline_rows = 0,
853 baseline_logical_bytes = 0,
854 dimension_rows = 0,
855 instance_rows = 0,
856 label_rows = 0,
857 daily_rows = 0,
858 label_tombstone_rows = 0,
859 instance_tombstone_rows = 0,
860 pruned_raw_rows = 0,
861 pruned_instance_rows = 0,
862 evicted_tombstones = 0,
863 writes_since_checkpoint = 0,
864 last_maintenance_epoch = 0,
865 last_checkpoint_epoch = 0,
866 oldest_retained_epoch = NULL,
867 raw_detail_complete = 1,
868 dimension_detail_complete = 1,
869 label_history_complete = 1,
870 maintenance_pending = 0,
871 clock_anomaly = 0,
872 spill_state = 'not_applicable',
873 checkpoint_state = 'not_due';",
874 )?;
875 initialize_empty_storage(connection)
876}
877
878pub(crate) fn migrate_legacy_usage(connection: &Connection) -> DbResult<()> {
880 let policy = TelemetryRetentionPolicy::default().validate()?;
881 let (project, _) = crate::project_identity::ensure_project_identity(connection)?;
882 ensure_overflow_dimension(connection, OVERFLOW_DIMENSION)?;
883 let mut statement = connection.prepare(
884 "SELECT session_id, command, path, query,
885 estimated_tokens_without_projectatlas,
886 estimated_tokens_with_projectatlas, estimated_tokens_saved,
887 token_savings_bucket, provider, model, tokenizer_backend,
888 accuracy, baseline_kind, confidence, calculation_trace,
889 accounting_layer, estimate_method, denominator_kind,
890 baseline_identity, baseline_fingerprint, dedupe_scope,
891 created_at, unixepoch(created_at)
892 FROM usage_events_legacy
893 ORDER BY session_id, id",
894 )?;
895 let mut rows = statement.query([])?;
896 let mut previous: Option<UsageInstanceId> = None;
897 while let Some(row) = rows.next()? {
898 let label = row.get::<_, String>(0)?;
899 let created_text = row.get::<_, String>(21)?;
900 let created_at = row
901 .get::<_, Option<i64>>(22)?
902 .ok_or_else(|| DbError::InvalidEnum {
903 field: "usage_events_legacy.created_at",
904 value: created_text,
905 })?;
906 let event = UsageEvent {
907 session_id: label.clone(),
908 command: row.get(1)?,
909 path: row.get(2)?,
910 query: row.get(3)?,
911 estimated_tokens_without_projectatlas: row
912 .get::<_, Option<i64>>(4)?
913 .map(|value| count_usize("estimated_tokens_without_projectatlas", value))
914 .transpose()?,
915 estimated_tokens_with_projectatlas: row
916 .get::<_, Option<i64>>(5)?
917 .map(|value| count_usize("estimated_tokens_with_projectatlas", value))
918 .transpose()?,
919 estimated_tokens_saved: row.get(6)?,
920 token_savings_bucket: row.get(7)?,
921 provider: row.get(8)?,
922 model: row.get(9)?,
923 tokenizer_backend: row.get(10)?,
924 accuracy: row.get(11)?,
925 baseline_kind: row.get(12)?,
926 confidence: row.get(13)?,
927 calculation_trace: row.get(14)?,
928 accounting_layer: row.get(15)?,
929 estimate_method: row.get(16)?,
930 denominator_kind: row.get(17)?,
931 baseline_identity: row.get(18)?,
932 baseline_fingerprint: row.get(19)?,
933 dedupe_scope: row.get(20)?,
934 };
935 let instance = migrated_instance_id(project, &label)?;
936 if previous != Some(instance) {
937 if let Some(previous) = previous {
938 seal_usage_instance_for_project(connection, project, previous, created_at)?;
939 }
940 previous = Some(instance);
941 }
942 let (event, detail_loss) = normalize_legacy_event(event, policy)?;
943 validate_event(&event, policy)?;
944 record_usage_at(
945 connection,
946 project,
947 instance,
948 UsageInstanceOwner::MigratedLegacy,
949 None,
950 false,
951 &event,
952 policy,
953 created_at,
954 false,
955 BaselineAdmission::SupportedUpgrade,
956 if detail_loss.dimension {
957 DimensionAdmission::Overflow
958 } else {
959 DimensionAdmission::Event
960 },
961 )?;
962 mark_legacy_detail_loss(connection, project, instance, &event, detail_loss)?;
963 }
964 drop(rows);
965 drop(statement);
966 if let Some(instance) = previous {
967 seal_usage_instance_for_project(connection, project, instance, now_epoch_seconds()?)?;
968 }
969 reconcile_retention_counters(connection)?;
970 converge_retention(connection, project, policy, now_epoch_seconds()?)
971}
972
973pub(crate) fn record_usage_for_project(
975 connection: &Connection,
976 project: ProjectInstanceId,
977 instance_id: UsageInstanceId,
978 owner: UsageInstanceOwner,
979 worktree_registration_id: Option<i64>,
980 event: &UsageEvent,
981 policy: TelemetryRetentionPolicy,
982 seal_after_record: bool,
983) -> DbResult<()> {
984 crate::project_identity::require_bound_project_identity(connection, project)?;
985 let policy = policy.validate()?;
986 validate_event(event, policy)?;
987 record_usage_at(
988 connection,
989 project,
990 instance_id,
991 owner,
992 worktree_registration_id,
993 true,
994 event,
995 policy,
996 now_epoch_seconds()?,
997 seal_after_record,
998 BaselineAdmission::BoundedRuntime,
999 DimensionAdmission::Event,
1000 )
1001}
1002
1003#[allow(clippy::too_many_arguments)]
1004fn record_usage_at(
1006 connection: &Connection,
1007 project: ProjectInstanceId,
1008 instance_id: UsageInstanceId,
1009 owner: UsageInstanceOwner,
1010 worktree_registration_id: Option<i64>,
1011 track_aggregate_revision: bool,
1012 event: &UsageEvent,
1013 policy: TelemetryRetentionPolicy,
1014 now: i64,
1015 seal_after_record: bool,
1016 baseline_admission: BaselineAdmission,
1017 dimension_admission: DimensionAdmission,
1018) -> DbResult<()> {
1019 expire_idle_instances(connection, project, instance_id, policy, now)?;
1020 let instance_exists = connection.query_row(
1021 "SELECT EXISTS(
1022 SELECT 1 FROM usage_instances
1023 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2
1024 )",
1025 params![
1026 project.as_bytes().as_slice(),
1027 instance_id.as_bytes().as_slice()
1028 ],
1029 |row| row.get::<_, i64>(0),
1030 )?;
1031 let reserve_instances = usize::from(instance_exists == 0);
1032 let (pruned_instances, instance_raw_rows) =
1033 prune_instances_once(connection, policy, now, reserve_instances)?;
1034 let instance_row_id = ensure_active_instance(
1035 connection,
1036 project,
1037 instance_id,
1038 owner,
1039 event_label(event),
1040 policy,
1041 now,
1042 )?;
1043 if track_aggregate_revision || worktree_registration_id.is_some() {
1044 bind_worktree_origin(connection, instance_row_id, worktree_registration_id)?;
1045 }
1046 ensure_label(connection, project, event_label(event), policy, now)?;
1047 let dimension = match dimension_admission {
1048 DimensionAdmission::Event => DimensionValues::from_event(event),
1049 DimensionAdmission::Overflow => DimensionValues::overflow(event.report_denominator_kind()),
1050 };
1051 let dimension_id = ensure_dimension(connection, &dimension, policy)?;
1052 let logical_bytes = logical_event_bytes(event, event_label(event))?;
1053 let delta = aggregate_delta(
1054 connection,
1055 instance_row_id,
1056 event,
1057 policy,
1058 baseline_admission,
1059 )?;
1060 insert_raw_event(
1061 connection,
1062 instance_row_id,
1063 dimension_id,
1064 event,
1065 now,
1066 logical_bytes,
1067 )?;
1068 prune_daily_once(connection, policy, now)?;
1069 apply_aggregates(
1070 connection,
1071 project,
1072 instance_row_id,
1073 dimension_id,
1074 now,
1075 delta,
1076 policy,
1077 )?;
1078 if let Some(registration_id) = worktree_registration_id {
1079 upsert_routed_worktree_aggregates(
1080 connection,
1081 registration_id,
1082 dimension_id,
1083 now,
1084 delta,
1085 policy,
1086 )?;
1087 }
1088 if track_aggregate_revision {
1089 increment_aggregate_revision(connection, project, delta.calls)?;
1090 }
1091 touch_instance(connection, instance_row_id, now, policy)?;
1092 if seal_after_record {
1093 seal_usage_instance_for_project(connection, project, instance_id, now)?;
1094 }
1095 let pruned_raw = prune_raw_once(connection, policy, now)?
1096 .checked_add(instance_raw_rows)
1097 .ok_or(DbError::TelemetryIntegerOverflow {
1098 field: "pruned_raw_rows",
1099 })?;
1100 let evicted_tombstones = prune_tombstones_once(connection, policy, now)?;
1101 prune_labels_once(connection, policy, now)?;
1102 refresh_retention_state(
1103 connection,
1104 policy,
1105 now,
1106 pruned_raw,
1107 pruned_instances,
1108 evicted_tombstones,
1109 1,
1110 )?;
1111 Ok(())
1112}
1113
1114pub(crate) fn seal_usage_instance(
1116 connection: &Connection,
1117 instance_id: UsageInstanceId,
1118) -> DbResult<()> {
1119 let project = current_project(connection)?;
1120 seal_usage_instance_for_project(connection, project, instance_id, now_epoch_seconds()?)
1121}
1122
1123fn seal_usage_instance_for_project(
1125 connection: &Connection,
1126 project: ProjectInstanceId,
1127 instance_id: UsageInstanceId,
1128 now: i64,
1129) -> DbResult<()> {
1130 let changed = connection.execute(
1131 "UPDATE usage_instances
1132 SET state = ?3, sealed_at_epoch = CASE
1133 WHEN last_seen_at_epoch > ?4 THEN last_seen_at_epoch ELSE ?4 END
1134 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2 AND state = ?5",
1135 params![
1136 project.as_bytes().as_slice(),
1137 instance_id.as_bytes().as_slice(),
1138 INSTANCE_SEALED,
1139 now,
1140 INSTANCE_ACTIVE,
1141 ],
1142 )?;
1143 if changed == 0 {
1144 return Err(DbError::TelemetryInstanceInactive);
1145 }
1146 let row_id = connection.query_row(
1147 "SELECT instance_row_id FROM usage_instances
1148 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2",
1149 params![
1150 project.as_bytes().as_slice(),
1151 instance_id.as_bytes().as_slice()
1152 ],
1153 |row| row.get::<_, i64>(0),
1154 )?;
1155 delete_instance_baselines(connection, row_id)?;
1156 Ok(())
1157}
1158
1159pub(crate) fn seal_project_usage_instances(
1161 connection: &Connection,
1162 project: ProjectInstanceId,
1163) -> DbResult<usize> {
1164 crate::project_identity::require_bound_project_identity(connection, project)?;
1165 let now = now_epoch_seconds()?;
1166 let project_bytes = project.as_bytes();
1167 let (baseline_rows, baseline_bytes) = connection.query_row(
1168 "SELECT COUNT(*), COALESCE(SUM(b.witness_logical_bytes), 0)
1169 FROM usage_instance_baselines AS b
1170 JOIN usage_instances AS i USING(instance_row_id)
1171 WHERE i.project_instance_id = ?1 AND i.state = ?2",
1172 params![project_bytes.as_slice(), INSTANCE_ACTIVE],
1173 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
1174 )?;
1175 connection.execute(
1176 "DELETE FROM usage_instance_baselines
1177 WHERE instance_row_id IN (
1178 SELECT instance_row_id FROM usage_instances
1179 WHERE project_instance_id = ?1 AND state = ?2
1180 )",
1181 params![project_bytes.as_slice(), INSTANCE_ACTIVE],
1182 )?;
1183 decrement_retention_counter(
1184 connection,
1185 RetentionCounter::BaselineRows,
1186 count_usize("project_baseline_rows", baseline_rows)?,
1187 )?;
1188 decrement_retention_counter(
1189 connection,
1190 RetentionCounter::BaselineLogicalBytes,
1191 count_usize("project_baseline_logical_bytes", baseline_bytes)?,
1192 )?;
1193 connection
1194 .execute(
1195 "UPDATE usage_instances
1196 SET state = ?2,
1197 sealed_at_epoch = CASE
1198 WHEN last_seen_at_epoch > ?3 THEN last_seen_at_epoch ELSE ?3 END
1199 WHERE project_instance_id = ?1 AND state = ?4",
1200 params![
1201 project_bytes.as_slice(),
1202 INSTANCE_SEALED,
1203 now,
1204 INSTANCE_ACTIVE,
1205 ],
1206 )
1207 .map_err(Into::into)
1208}
1209
1210#[derive(Clone, Copy)]
1212enum MaintenanceRoot<'a> {
1213 #[cfg(test)]
1215 Display(Option<&'a str>),
1216 Native(Option<&'a projectatlas_core::CanonicalProjectRoot>),
1218}
1219
1220#[cfg(test)]
1222pub(crate) fn maintain_after_commit_for_project(
1223 connection: &Connection,
1224 expected_root: Option<&str>,
1225 project: ProjectInstanceId,
1226 policy: TelemetryRetentionPolicy,
1227) -> DbResult<()> {
1228 maintain_after_commit_for_project_with_checkpoint_binding(
1229 connection,
1230 MaintenanceRoot::Display(expected_root),
1231 project,
1232 policy,
1233 |connection| Ok(passive_checkpoint_state(connection)),
1234 )
1235}
1236
1237pub(crate) fn maintain_after_commit_for_native_project(
1239 connection: &Connection,
1240 expected_root: Option<&projectatlas_core::CanonicalProjectRoot>,
1241 project: ProjectInstanceId,
1242 policy: TelemetryRetentionPolicy,
1243) -> DbResult<()> {
1244 maintain_after_commit_for_project_with_checkpoint_binding(
1245 connection,
1246 MaintenanceRoot::Native(expected_root),
1247 project,
1248 policy,
1249 |connection| Ok(passive_checkpoint_state(connection)),
1250 )
1251}
1252
1253#[cfg(test)]
1255fn maintain_after_commit_for_project_with_checkpoint(
1256 connection: &Connection,
1257 expected_root: Option<&str>,
1258 project: ProjectInstanceId,
1259 policy: TelemetryRetentionPolicy,
1260 checkpoint: impl FnOnce(&Connection) -> DbResult<TelemetryCheckpointState>,
1261) -> DbResult<()> {
1262 maintain_after_commit_for_project_with_checkpoint_binding(
1263 connection,
1264 MaintenanceRoot::Display(expected_root),
1265 project,
1266 policy,
1267 checkpoint,
1268 )
1269}
1270
1271fn maintain_after_commit_for_project_with_checkpoint_binding(
1273 connection: &Connection,
1274 expected_root: MaintenanceRoot<'_>,
1275 project: ProjectInstanceId,
1276 policy: TelemetryRetentionPolicy,
1277 checkpoint: impl FnOnce(&Connection) -> DbResult<TelemetryCheckpointState>,
1278) -> DbResult<()> {
1279 crate::project_identity::require_bound_project_identity(connection, project)?;
1280 let policy = policy.validate()?;
1281 let checkpoint_start_writes = count_usize(
1282 "writes_since_checkpoint",
1283 connection.query_row(
1284 "SELECT writes_since_checkpoint FROM usage_retention_state WHERE singleton = 1",
1285 [],
1286 |row| row.get::<_, i64>(0),
1287 )?,
1288 )?;
1289 if checkpoint_start_writes < policy.checkpoint_write_interval {
1290 return Ok(());
1291 }
1292 match expected_root {
1293 #[cfg(test)]
1294 MaintenanceRoot::Display(root) => {
1295 crate::schema::validate_active_binding(connection, root, Some(project))?;
1296 }
1297 MaintenanceRoot::Native(root) => {
1298 crate::schema::validate_active_native_binding(connection, root, Some(project))?;
1299 }
1300 }
1301 let checkpoint_data_version =
1302 connection.query_row("PRAGMA data_version", [], |row| row.get::<_, i64>(0))?;
1303 let state = checkpoint(connection)?;
1304 with_maintenance_write_transaction(connection, expected_root, project, |transaction| {
1305 let current_data_version =
1306 transaction.query_row("PRAGMA data_version", [], |row| row.get::<_, i64>(0))?;
1307 let current_writes = count_usize(
1308 "writes_since_checkpoint",
1309 transaction.query_row(
1310 "SELECT writes_since_checkpoint
1311 FROM usage_retention_state WHERE singleton = 1",
1312 [],
1313 |row| row.get::<_, i64>(0),
1314 )?,
1315 )?;
1316 let retained_writes = writes_after_checkpoint_attempt(
1317 state,
1318 checkpoint_start_writes,
1319 current_writes,
1320 current_data_version == checkpoint_data_version,
1321 );
1322 let now = now_epoch_seconds()?;
1323 let (raw_rows, raw_bytes, old_raw) = raw_pressure(transaction, policy, now)?;
1324 let retention_pending = raw_rows > policy.max_raw_rows
1325 || raw_bytes > policy.max_raw_logical_bytes
1326 || old_raw > 0
1327 || retention_counter(transaction, RetentionCounter::InstanceRows)?
1328 > policy.max_retained_instances
1329 || retention_counter(transaction, RetentionCounter::LabelRows)?
1330 > policy.max_retained_labels
1331 || retention_counter(transaction, RetentionCounter::DailyRows)? > policy.max_daily_rows
1332 || retention_counter(transaction, RetentionCounter::LabelTombstoneRows)?
1333 > policy.max_label_tombstones
1334 || retention_counter(transaction, RetentionCounter::InstanceTombstoneRows)?
1335 > policy.max_instance_tombstones
1336 || aged_maintenance_pending(transaction, policy, now)?;
1337 let completed = state == TelemetryCheckpointState::Completed;
1338 transaction.execute(
1339 "UPDATE usage_retention_state
1340 SET writes_since_checkpoint = ?1,
1341 last_checkpoint_epoch = ?2,
1342 checkpoint_state = ?3,
1343 maintenance_pending = ?4
1344 WHERE singleton = 1",
1345 params![
1346 to_i64("writes_since_checkpoint", retained_writes)?,
1347 now,
1348 state.as_str(),
1349 i64::from(
1350 !completed
1351 || retention_pending
1352 || retained_writes >= policy.checkpoint_write_interval
1353 ),
1354 ],
1355 )?;
1356 Ok(())
1357 })
1358}
1359
1360fn with_maintenance_write_transaction<T>(
1362 connection: &Connection,
1363 expected_root: MaintenanceRoot<'_>,
1364 project: ProjectInstanceId,
1365 operation: impl FnOnce(&Connection) -> DbResult<T>,
1366) -> DbResult<T> {
1367 match expected_root {
1368 #[cfg(test)]
1369 MaintenanceRoot::Display(root) => {
1370 crate::with_validated_write_transaction(connection, root, Some(project), operation)
1371 }
1372 MaintenanceRoot::Native(root) => crate::with_validated_native_write_transaction(
1373 connection,
1374 root,
1375 Some(project),
1376 operation,
1377 ),
1378 }
1379}
1380
1381fn passive_checkpoint_state(connection: &Connection) -> TelemetryCheckpointState {
1383 let result = connection.query_row("PRAGMA wal_checkpoint(PASSIVE)", [], |row| {
1384 Ok((
1385 row.get::<_, i64>(0)?,
1386 row.get::<_, i64>(1)?,
1387 row.get::<_, i64>(2)?,
1388 ))
1389 });
1390 match result {
1391 Ok((0, log_frames, checkpointed_frames)) if checkpointed_frames == log_frames => {
1392 TelemetryCheckpointState::Completed
1393 }
1394 Ok(_) => TelemetryCheckpointState::Busy,
1395 Err(_) => TelemetryCheckpointState::Error,
1396 }
1397}
1398
1399const fn writes_after_checkpoint_attempt(
1401 state: TelemetryCheckpointState,
1402 checkpoint_start_writes: usize,
1403 current_writes: usize,
1404 data_version_unchanged: bool,
1405) -> usize {
1406 if matches!(state, TelemetryCheckpointState::Completed) && data_version_unchanged {
1407 current_writes.saturating_sub(checkpoint_start_writes)
1408 } else {
1409 current_writes
1410 }
1411}
1412
1413pub(crate) fn retention_state(connection: &Connection) -> DbResult<TelemetryRetentionState> {
1415 let project = current_project(connection)?;
1416 retention_state_for_project(connection, project)
1417}
1418
1419fn retention_state_for_project(
1421 connection: &Connection,
1422 project: ProjectInstanceId,
1423) -> DbResult<TelemetryRetentionState> {
1424 crate::project_identity::require_bound_project_identity(connection, project)?;
1425 let row = connection.query_row(
1426 "SELECT policy_version, logical_byte_version, raw_rows, raw_logical_bytes,
1427 baseline_rows, baseline_logical_bytes, dimension_rows, instance_rows,
1428 daily_rows, label_tombstone_rows, instance_tombstone_rows,
1429 pruned_raw_rows, pruned_instance_rows, evicted_tombstones,
1430 maintenance_pending, clock_anomaly, spill_state, checkpoint_state
1431 FROM usage_retention_state WHERE singleton = 1",
1432 [],
1433 |row| {
1434 Ok((
1435 row.get::<_, i64>(0)?,
1436 row.get::<_, i64>(1)?,
1437 row.get::<_, i64>(2)?,
1438 row.get::<_, i64>(3)?,
1439 row.get::<_, i64>(4)?,
1440 row.get::<_, i64>(5)?,
1441 row.get::<_, i64>(6)?,
1442 row.get::<_, i64>(7)?,
1443 row.get::<_, i64>(8)?,
1444 row.get::<_, i64>(9)?,
1445 row.get::<_, i64>(10)?,
1446 row.get::<_, i64>(11)?,
1447 row.get::<_, i64>(12)?,
1448 row.get::<_, i64>(13)?,
1449 row.get::<_, i64>(14)?,
1450 row.get::<_, i64>(15)?,
1451 row.get::<_, String>(16)?,
1452 row.get::<_, String>(17)?,
1453 ))
1454 },
1455 )?;
1456 if row.16 != "not_applicable" {
1457 return Err(DbError::InvalidEnum {
1458 field: "usage_retention_state.spill_state",
1459 value: row.16,
1460 });
1461 }
1462 let policy = TelemetryRetentionPolicy::default().validate()?;
1463 let lifecycle = connection.query_row(
1464 "SELECT writes_since_checkpoint, last_checkpoint_epoch, oldest_retained_epoch
1465 FROM usage_retention_state WHERE singleton = 1",
1466 [],
1467 |row| {
1468 Ok((
1469 row.get::<_, i64>(0)?,
1470 row.get::<_, i64>(1)?,
1471 row.get::<_, Option<i64>>(2)?,
1472 ))
1473 },
1474 )?;
1475 let active_instances = connection.query_row(
1476 "SELECT COUNT(*) FROM usage_instances
1477 WHERE project_instance_id = ?1 AND state = ?2",
1478 params![project.as_bytes().as_slice(), INSTANCE_ACTIVE],
1479 |row| row.get::<_, i64>(0),
1480 )?;
1481 let retained_labels = retention_counter(connection, RetentionCounter::LabelRows)?;
1482 let journal_mode =
1483 connection.query_row("PRAGMA journal_mode", [], |row| row.get::<_, String>(0))?;
1484 let synchronous = connection.query_row("PRAGMA synchronous", [], |row| row.get::<_, i64>(0))?;
1485 let busy_timeout =
1486 connection.query_row("PRAGMA busy_timeout", [], |row| row.get::<_, i64>(0))?;
1487 let statistics = connection.query_row(
1488 "SELECT EXISTS(
1489 SELECT 1 FROM sqlite_schema WHERE type = 'table' AND name = 'sqlite_stat1'
1490 )",
1491 [],
1492 |row| row.get::<_, i64>(0),
1493 )?;
1494 let checkpoint_state = TelemetryCheckpointState::from_str(&row.17)?;
1495 let wal_autocheckpoint =
1496 connection.query_row("PRAGMA wal_autocheckpoint", [], |row| row.get::<_, i64>(0))?;
1497 let normal_busy_timeout_ms =
1498 duration_millis("normal_busy_timeout_ms", crate::SQLITE_BUSY_TIMEOUT)?;
1499 let telemetry_busy_timeout_ms = duration_millis(
1500 "telemetry_busy_timeout_ms",
1501 crate::SQLITE_TELEMETRY_BUSY_TIMEOUT,
1502 )?;
1503 Ok(TelemetryRetentionState {
1504 policy_version: count_u32("policy_version", row.0)?,
1505 logical_byte_version: count_u32("logical_byte_version", row.1)?,
1506 raw_rows: count_usize("raw_rows", row.2)?,
1507 max_raw_rows: policy.max_raw_rows,
1508 max_raw_age_seconds: policy.max_raw_age_seconds,
1509 raw_logical_bytes: count_usize("raw_logical_bytes", row.3)?,
1510 max_raw_logical_bytes: policy.max_raw_logical_bytes,
1511 baseline_rows: count_usize("baseline_rows", row.4)?,
1512 max_baselines_per_instance: policy.max_baselines_per_instance,
1513 max_active_baseline_rows: policy.max_active_baseline_rows,
1514 baseline_logical_bytes: count_usize("baseline_logical_bytes", row.5)?,
1515 max_baseline_logical_bytes: policy.max_baseline_logical_bytes,
1516 dimension_rows: count_usize("dimension_rows", row.6)?,
1517 max_dimensions: policy.max_dimensions,
1518 instance_rows: count_usize("instance_rows", row.7)?,
1519 active_instance_rows: count_usize("active_instance_rows", active_instances)?,
1520 max_active_instances: policy.max_active_instances,
1521 max_retained_instances: policy.max_retained_instances,
1522 retained_label_rows: retained_labels,
1523 max_retained_labels: policy.max_retained_labels,
1524 daily_rows: count_usize("daily_rows", row.8)?,
1525 max_daily_rows: policy.max_daily_rows,
1526 retained_trend_days: policy.retained_trend_days,
1527 label_tombstone_rows: count_usize("label_tombstone_rows", row.9)?,
1528 max_label_tombstones: policy.max_label_tombstones,
1529 instance_tombstone_rows: count_usize("instance_tombstone_rows", row.10)?,
1530 max_instance_tombstones: policy.max_instance_tombstones,
1531 pruned_raw_rows: count_usize("pruned_raw_rows", row.11)?,
1532 pruned_instance_rows: count_usize("pruned_instance_rows", row.12)?,
1533 evicted_tombstones: count_usize("evicted_tombstones", row.13)?,
1534 maintenance_pending: bool_from_sql("maintenance_pending", row.14)?,
1535 prune_batch_rows: policy.prune_batch_rows,
1536 writes_since_checkpoint: count_usize("writes_since_checkpoint", lifecycle.0)?,
1537 checkpoint_write_interval: policy.checkpoint_write_interval,
1538 last_checkpoint_epoch: count_u64("last_checkpoint_epoch", lifecycle.1)?,
1539 oldest_retained_epoch: lifecycle
1540 .2
1541 .map(|value| count_u64("oldest_retained_epoch", value))
1542 .transpose()?,
1543 clock_anomaly: bool_from_sql("clock_anomaly", row.15)?,
1544 spill_cleanup: SpillCleanupState::NotApplicable,
1545 checkpoint_state,
1546 wal_autocheckpoint_pages: count_usize("wal_autocheckpoint", wal_autocheckpoint)?,
1547 freelist_pages: count_usize(
1548 "freelist_pages",
1549 pragma_count(connection, "freelist_count")?,
1550 )?,
1551 page_count: count_usize("page_count", pragma_count(connection, "page_count")?)?,
1552 page_size: count_usize("page_size", pragma_count(connection, "page_size")?)?,
1553 journal_mode,
1554 synchronous_mode: synchronous_mode(synchronous)?.to_string(),
1555 connection_busy_timeout_ms: count_u64("busy_timeout", busy_timeout)?,
1556 normal_busy_timeout_ms,
1557 telemetry_busy_timeout_ms,
1558 statistics_policy: PlannerStatisticsPolicy::NotConfigured,
1559 statistics_state: if statistics == 0 {
1560 PlannerStatisticsState::NotInitialized
1561 } else {
1562 PlannerStatisticsState::Available
1563 },
1564 })
1565}
1566
1567pub(crate) fn usage_events(
1569 connection: &Connection,
1570 caller_label: Option<&str>,
1571) -> DbResult<Vec<UsageEvent>> {
1572 let project = current_project(connection)?;
1573 usage_events_for_project(connection, project, caller_label)
1574}
1575
1576fn usage_events_for_project(
1578 connection: &Connection,
1579 project: ProjectInstanceId,
1580 caller_label: Option<&str>,
1581) -> DbResult<Vec<UsageEvent>> {
1582 crate::project_identity::require_bound_project_identity(connection, project)?;
1583 let sql = if caller_label.is_some() {
1584 raw_event_select("AND i.caller_label = ?2")
1585 } else {
1586 raw_event_select("")
1587 };
1588 let mut statement = connection.prepare(&sql)?;
1589 let mut rows = if let Some(label) = caller_label {
1590 statement.query(params![project.as_bytes().as_slice(), label])?
1591 } else {
1592 statement.query([project.as_bytes().as_slice()])?
1593 };
1594 let mut events = Vec::new();
1595 while let Some(row) = rows.next()? {
1596 events.push(map_usage_event(row)?);
1597 }
1598 Ok(events)
1599}
1600
1601pub(crate) fn export_worktree_usage_snapshot(
1603 connection: &Connection,
1604) -> DbResult<WorktreeUsageSnapshot> {
1605 let owned_snapshot = connection
1606 .is_autocommit()
1607 .then(|| {
1608 rusqlite::Transaction::new_unchecked(
1609 connection,
1610 rusqlite::TransactionBehavior::Deferred,
1611 )
1612 })
1613 .transpose()?;
1614 let connection = owned_snapshot.as_deref().unwrap_or(connection);
1615 let project = current_project(connection)?;
1616 crate::project_identity::require_bound_project_identity(connection, project)?;
1617 let policy = TelemetryRetentionPolicy::default().validate()?;
1618 let revision = connection
1619 .query_row(
1620 "SELECT revision FROM usage_aggregate_revisions
1621 WHERE project_instance_id = ?1",
1622 [project.as_bytes().as_slice()],
1623 |row| row.get::<_, i64>(0),
1624 )
1625 .optional()?
1626 .unwrap_or(0);
1627 let revision =
1628 u64::try_from(revision).map_err(|_source| DbError::TelemetryIntegerOverflow {
1629 field: "usage_aggregate_revisions.revision",
1630 })?;
1631 let maximum_rows = policy
1632 .max_daily_rows
1633 .checked_add(policy.max_dimensions)
1634 .and_then(|value| value.checked_add(2))
1635 .ok_or(DbError::TelemetryIntegerOverflow {
1636 field: "worktree_snapshot_rows",
1637 })?;
1638 let query_limit = to_i64("worktree_snapshot_rows", maximum_rows.saturating_add(1))?;
1639 let mut statement = connection.prepare(
1640 "SELECT -1 AS day_epoch, dimension_id,
1641 calls, estimated_without, estimated_with, observed_without,
1642 observed_with, modeled_without, modeled_with,
1643 deduped_modeled_without, deduped_modeled_with, repeated_baselines,
1644 observed_file_read_replacements, modeled_file_reads_avoided
1645 FROM usage_global_aggregates
1646 WHERE project_instance_id = ?1
1647 UNION ALL
1648 SELECT day_epoch, dimension_id,
1649 calls, estimated_without, estimated_with, observed_without,
1650 observed_with, modeled_without, modeled_with,
1651 deduped_modeled_without, deduped_modeled_with, repeated_baselines,
1652 observed_file_read_replacements, modeled_file_reads_avoided
1653 FROM usage_daily_aggregates
1654 WHERE project_instance_id = ?1
1655 ORDER BY day_epoch, dimension_id
1656 LIMIT ?2",
1657 )?;
1658 let rows = statement
1659 .query_map(params![project.as_bytes().as_slice(), query_limit], |row| {
1660 Ok(WorktreeUsageSnapshotRow {
1661 day_epoch: row.get(0)?,
1662 dimension_id: row.get(1)?,
1663 counters: read_counters_offset(row, 2).map_err(|error| {
1664 rusqlite::Error::FromSqlConversionFailure(
1665 2,
1666 rusqlite::types::Type::Integer,
1667 Box::new(error),
1668 )
1669 })?,
1670 })
1671 })?
1672 .collect::<Result<Vec<_>, _>>()?;
1673 if rows.len() > maximum_rows {
1674 return Err(DbError::WorktreeTelemetrySnapshotLimit {
1675 resource: "rows",
1676 limit: maximum_rows,
1677 observed: rows.len(),
1678 });
1679 }
1680 let referenced_dimensions = rows
1681 .iter()
1682 .map(|row| row.dimension_id)
1683 .collect::<BTreeSet<_>>();
1684 let mut dimensions = BTreeMap::new();
1685 let mut dimension_statement = connection.prepare(
1686 "SELECT dimension_id,
1687 token_savings_bucket, provider, model, tokenizer_backend,
1688 accuracy, baseline_kind, confidence, accounting_layer,
1689 estimate_method, denominator_kind, dedupe_scope, overflow
1690 FROM usage_bucket_dimensions
1691 ORDER BY dimension_id",
1692 )?;
1693 let mut dimension_rows = dimension_statement.query([])?;
1694 while let Some(row) = dimension_rows.next()? {
1695 let dimension_id = row.get::<_, i64>(0)?;
1696 if referenced_dimensions.contains(&dimension_id) {
1697 dimensions.insert(dimension_id, read_dimension(row, 1)?);
1698 }
1699 }
1700 drop(dimension_rows);
1701 drop(dimension_statement);
1702 drop(statement);
1703 if dimensions.len() != referenced_dimensions.len() {
1704 return Err(DbError::WorktreeRegistrationRow {
1705 reason: "aggregate snapshot references a missing dimension",
1706 });
1707 }
1708 let logical_bytes = validate_worktree_usage_snapshot(&dimensions, &rows, policy)?;
1709 let snapshot = WorktreeUsageSnapshot {
1710 project_instance_id: project,
1711 revision,
1712 dimensions,
1713 rows,
1714 logical_bytes,
1715 };
1716 if let Some(snapshot) = owned_snapshot {
1717 snapshot.commit()?;
1718 }
1719 Ok(snapshot)
1720}
1721
1722pub(crate) fn synchronize_worktree_usage_snapshot(
1724 connection: &Connection,
1725 registration_id: i64,
1726 snapshot: &WorktreeUsageSnapshot,
1727) -> DbResult<WorktreeUsageSyncState> {
1728 let policy = TelemetryRetentionPolicy::default().validate()?;
1729 let (project_bytes, accepted_revision) = connection.query_row(
1730 "SELECT project_instance_id, accepted_telemetry_revision
1731 FROM worktree_registrations
1732 WHERE registration_id = ?1 AND state = 'active'",
1733 [registration_id],
1734 |row| Ok((row.get::<_, Option<Vec<u8>>>(0)?, row.get::<_, i64>(1)?)),
1735 )?;
1736 let Some(project_bytes) = project_bytes else {
1737 return Err(DbError::WorktreeTelemetryProjectMismatch { registration_id });
1738 };
1739 let project_bytes: [u8; 16] =
1740 project_bytes
1741 .try_into()
1742 .map_err(|value: Vec<u8>| DbError::InvalidBlobLength {
1743 field: "worktree_registrations.project_instance_id",
1744 expected: 16,
1745 found: value.len(),
1746 })?;
1747 let project = ProjectInstanceId::from_bytes(project_bytes).map_err(DbError::from)?;
1748 if project != snapshot.project_instance_id {
1749 return Err(DbError::WorktreeTelemetryProjectMismatch { registration_id });
1750 }
1751 let accepted_revision =
1752 u64::try_from(accepted_revision).map_err(|_source| DbError::WorktreeRegistrationRow {
1753 reason: "negative accepted telemetry revision",
1754 })?;
1755 if snapshot.revision <= accepted_revision {
1756 return Ok(WorktreeUsageSyncState::Current);
1757 }
1758 let logical_bytes =
1759 validate_worktree_usage_snapshot(&snapshot.dimensions, &snapshot.rows, policy)?;
1760 if logical_bytes != snapshot.logical_bytes {
1761 return Err(DbError::WorktreeRegistrationRow {
1762 reason: "aggregate snapshot logical-byte contract changed",
1763 });
1764 }
1765 let retained_trend_seconds = policy
1766 .retained_trend_days
1767 .checked_mul(SECONDS_PER_DAY as u64)
1768 .ok_or(DbError::TelemetryIntegerOverflow {
1769 field: "retained_trend_seconds",
1770 })?;
1771 let retained_daily_cutoff = epoch_cutoff(
1772 now_epoch_seconds()?,
1773 retained_trend_seconds,
1774 "retained_trend_seconds",
1775 )?;
1776 connection.execute(
1777 "DELETE FROM worktree_usage_aggregates
1778 WHERE day_epoch >= 0 AND day_epoch < ?1 AND source_kind = ?2",
1779 params![retained_daily_cutoff, WORKTREE_USAGE_SYNCHRONIZED],
1780 )?;
1781 let mut target_dimensions = BTreeMap::new();
1782 for (source_id, dimension) in &snapshot.dimensions {
1783 target_dimensions.insert(*source_id, ensure_dimension(connection, dimension, policy)?);
1784 }
1785 let mut incoming = BTreeMap::<(i64, i64), AggregateCounters>::new();
1786 for row in snapshot
1787 .rows
1788 .iter()
1789 .filter(|row| row.day_epoch == -1 || row.day_epoch >= retained_daily_cutoff)
1790 {
1791 let dimension_id =
1792 *target_dimensions
1793 .get(&row.dimension_id)
1794 .ok_or(DbError::WorktreeRegistrationRow {
1795 reason: "aggregate snapshot dimension mapping is incomplete",
1796 })?;
1797 let total = incoming.entry((row.day_epoch, dimension_id)).or_default();
1798 *total = total.checked_add(row.counters)?;
1799 }
1800 let accepted_lifetime = {
1801 let mut statement = connection.prepare_cached(
1802 "SELECT a.dimension_id,
1803 a.calls, a.estimated_without, a.estimated_with, a.observed_without,
1804 a.observed_with, a.modeled_without, a.modeled_with,
1805 a.deduped_modeled_without, a.deduped_modeled_with, a.repeated_baselines,
1806 a.observed_file_read_replacements, a.modeled_file_reads_avoided
1807 FROM worktree_usage_aggregates AS a
1808 INDEXED BY idx_worktree_usage_aggregates_day_registration
1809 WHERE a.registration_id = ?1 AND a.source_kind = ?2 AND a.day_epoch = -1
1810 ORDER BY a.dimension_id",
1811 )?;
1812 let mut rows = statement.query(params![registration_id, WORKTREE_USAGE_SYNCHRONIZED])?;
1813 let mut totals = BTreeMap::<i64, AggregateCounters>::new();
1814 while let Some(row) = rows.next()? {
1815 let total = totals.entry(row.get(0)?).or_default();
1816 *total = total.checked_add(read_counters_offset(row, 1)?)?;
1817 }
1818 totals
1819 };
1820 if accepted_lifetime.iter().any(|(dimension_id, accepted)| {
1821 !incoming
1822 .get(&(-1, *dimension_id))
1823 .is_some_and(|incoming| incoming.contains(*accepted))
1824 }) {
1825 return Err(DbError::WorktreeRegistrationRow {
1826 reason: "aggregate snapshot reduces accepted lifetime totals",
1827 });
1828 }
1829 let accepted_daily = {
1830 let mut statement = connection.prepare_cached(
1831 "SELECT a.day_epoch, a.dimension_id,
1832 a.calls, a.estimated_without, a.estimated_with, a.observed_without,
1833 a.observed_with, a.modeled_without, a.modeled_with,
1834 a.deduped_modeled_without, a.deduped_modeled_with,
1835 a.repeated_baselines, a.observed_file_read_replacements,
1836 a.modeled_file_reads_avoided
1837 FROM worktree_usage_aggregates AS a
1838 WHERE a.registration_id = ?1 AND a.source_kind = ?2
1839 AND a.day_epoch >= ?3
1840 ORDER BY a.day_epoch, a.dimension_id",
1841 )?;
1842 let mut rows = statement.query(params![
1843 registration_id,
1844 WORKTREE_USAGE_SYNCHRONIZED,
1845 retained_daily_cutoff
1846 ])?;
1847 let mut totals = BTreeMap::<(i64, i64), AggregateCounters>::new();
1848 while let Some(row) = rows.next()? {
1849 let key = (row.get(0)?, row.get(1)?);
1850 let total = totals.entry(key).or_default();
1851 *total = total.checked_add(read_counters_offset(row, 2)?)?;
1852 }
1853 totals
1854 };
1855 if accepted_daily.iter().any(|(key, accepted)| {
1856 !incoming
1857 .get(key)
1858 .is_some_and(|incoming| incoming.contains(*accepted))
1859 }) {
1860 return Err(DbError::WorktreeRegistrationRow {
1861 reason: "aggregate snapshot reduces accepted retained daily totals",
1862 });
1863 }
1864 let incoming_daily_rows = incoming
1865 .keys()
1866 .filter(|(day_epoch, _dimension_id)| *day_epoch >= 0)
1867 .count();
1868 let other_daily_rows = connection.query_row(
1869 "SELECT COUNT(*) FROM worktree_usage_aggregates
1870 WHERE day_epoch >= 0
1871 AND NOT (registration_id = ?1 AND source_kind = ?2)",
1872 params![registration_id, WORKTREE_USAGE_SYNCHRONIZED],
1873 |row| row.get::<_, i64>(0),
1874 )?;
1875 let projected_daily_rows = count_usize("worktree_daily_rows", other_daily_rows)?
1876 .checked_add(incoming_daily_rows)
1877 .ok_or(DbError::TelemetryIntegerOverflow {
1878 field: "worktree_daily_rows",
1879 })?;
1880 if projected_daily_rows > policy.max_daily_rows {
1881 return Err(DbError::WorktreeTelemetrySnapshotLimit {
1882 resource: "daily_rows",
1883 limit: policy.max_daily_rows,
1884 observed: projected_daily_rows,
1885 });
1886 }
1887 connection.execute(
1888 "DELETE FROM worktree_usage_aggregates
1889 WHERE registration_id = ?1 AND source_kind = ?2",
1890 params![registration_id, WORKTREE_USAGE_SYNCHRONIZED],
1891 )?;
1892 let mut insert = connection.prepare_cached(
1893 "INSERT INTO worktree_usage_aggregates(
1894 registration_id, source_kind, day_epoch, dimension_id,
1895 calls, estimated_without, estimated_with, observed_without,
1896 observed_with, modeled_without, modeled_with,
1897 deduped_modeled_without, deduped_modeled_with, repeated_baselines,
1898 observed_file_read_replacements, modeled_file_reads_avoided
1899 ) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
1900 )?;
1901 for ((day_epoch, dimension_id), counters) in incoming {
1902 insert.execute(worktree_aggregate_params!(
1903 registration_id,
1904 WORKTREE_USAGE_SYNCHRONIZED,
1905 day_epoch,
1906 dimension_id,
1907 counters
1908 ))?;
1909 }
1910 let revision =
1911 i64::try_from(snapshot.revision).map_err(|_source| DbError::TelemetryIntegerOverflow {
1912 field: "usage_aggregate_revisions.revision",
1913 })?;
1914 let updated = connection.execute(
1915 "UPDATE worktree_registrations
1916 SET accepted_telemetry_revision = ?2
1917 WHERE registration_id = ?1 AND state = 'active'
1918 AND accepted_telemetry_revision < ?2",
1919 params![registration_id, revision],
1920 )?;
1921 if updated != 1 {
1922 return Err(DbError::WorktreeRegistrationRow {
1923 reason: "aggregate snapshot revision changed during synchronization",
1924 });
1925 }
1926 Ok(WorktreeUsageSyncState::Synchronized)
1927}
1928
1929fn validate_worktree_usage_snapshot(
1931 dimensions: &BTreeMap<i64, DimensionValues>,
1932 rows: &[WorktreeUsageSnapshotRow],
1933 policy: TelemetryRetentionPolicy,
1934) -> DbResult<usize> {
1935 let maximum_rows = policy
1936 .max_daily_rows
1937 .checked_add(policy.max_dimensions)
1938 .and_then(|value| value.checked_add(2))
1939 .ok_or(DbError::TelemetryIntegerOverflow {
1940 field: "worktree_snapshot_rows",
1941 })?;
1942 if rows.len() > maximum_rows {
1943 return Err(DbError::WorktreeTelemetrySnapshotLimit {
1944 resource: "rows",
1945 limit: maximum_rows,
1946 observed: rows.len(),
1947 });
1948 }
1949 let mut logical_bytes = 24usize;
1950 for dimension in dimensions.values() {
1951 validate_dimension_values(dimension, policy.max_dimension_bytes)?;
1952 logical_bytes = logical_bytes
1953 .checked_add(9)
1954 .and_then(|value| value.checked_add(dimension_logical_bytes(dimension)))
1955 .ok_or(DbError::TelemetryIntegerOverflow {
1956 field: "worktree_snapshot_bytes",
1957 })?;
1958 }
1959 for row in rows {
1960 if row.day_epoch < -1 || !dimensions.contains_key(&row.dimension_id) {
1961 return Err(DbError::WorktreeRegistrationRow {
1962 reason: "aggregate snapshot row has an invalid day or dimension",
1963 });
1964 }
1965 row.counters.validate_nonnegative()?;
1966 logical_bytes =
1967 logical_bytes
1968 .checked_add(112)
1969 .ok_or(DbError::TelemetryIntegerOverflow {
1970 field: "worktree_snapshot_bytes",
1971 })?;
1972 }
1973 if logical_bytes > MAX_WORKTREE_USAGE_SNAPSHOT_BYTES {
1974 return Err(DbError::WorktreeTelemetrySnapshotLimit {
1975 resource: "logical_bytes",
1976 limit: MAX_WORKTREE_USAGE_SNAPSHOT_BYTES,
1977 observed: logical_bytes,
1978 });
1979 }
1980 Ok(logical_bytes)
1981}
1982
1983fn validate_dimension_values(dimension: &DimensionValues, maximum_bytes: usize) -> DbResult<()> {
1985 for (field, value) in [
1986 (
1987 "token_savings_bucket",
1988 dimension.token_savings_bucket.as_str(),
1989 ),
1990 ("provider", dimension.provider.as_str()),
1991 ("model", dimension.model.as_str()),
1992 ("tokenizer_backend", dimension.tokenizer_backend.as_str()),
1993 ("accuracy", dimension.accuracy.as_str()),
1994 ("baseline_kind", dimension.baseline_kind.as_str()),
1995 ("confidence", dimension.confidence.as_str()),
1996 ("accounting_layer", dimension.accounting_layer.as_str()),
1997 ("estimate_method", dimension.estimate_method.as_str()),
1998 ("denominator_kind", dimension.denominator_kind.as_str()),
1999 ("dedupe_scope", dimension.dedupe_scope.as_str()),
2000 ] {
2001 validate_required_text(field, value, maximum_bytes)?;
2002 }
2003 Ok(())
2004}
2005
2006fn dimension_logical_bytes(dimension: &DimensionValues) -> usize {
2008 [
2009 &dimension.token_savings_bucket,
2010 &dimension.provider,
2011 &dimension.model,
2012 &dimension.tokenizer_backend,
2013 &dimension.accuracy,
2014 &dimension.baseline_kind,
2015 &dimension.confidence,
2016 &dimension.accounting_layer,
2017 &dimension.estimate_method,
2018 &dimension.denominator_kind,
2019 &dimension.dedupe_scope,
2020 ]
2021 .iter()
2022 .map(|value| value.len())
2023 .sum()
2024}
2025
2026pub(crate) fn token_overview(
2028 connection: &Connection,
2029 caller_label: Option<&str>,
2030) -> DbResult<TokenOverview> {
2031 let project = current_project(connection)?;
2032 token_overview_for_project(connection, project, caller_label)
2033}
2034
2035pub(crate) fn repository_token_overview(connection: &Connection) -> DbResult<TokenOverview> {
2037 let project = current_project(connection)?;
2038 crate::project_identity::require_bound_project_identity(connection, project)?;
2039 let mut aggregates = load_overview_aggregates(connection, project, None)?;
2040 let worktree_aggregates = load_worktree_overview_aggregates(connection, None, true)?;
2041 let has_worktree_aggregates = worktree_aggregates_exist(connection)?;
2042 aggregates.extend(worktree_aggregates);
2043 let (buckets, totals, average_policy_complete) = aggregate_report_rows(aggregates)?;
2044 let mut overview = TokenOverview::from_buckets(buckets);
2045 overview.apply_accounting_totals(totals);
2046 if !average_policy_complete {
2047 overview.average_policy.evidence = TOKEN_AVERAGE_POLICY_OVERFLOW_EVIDENCE.to_string();
2048 }
2049 let native_detail = detail_availability(connection, project, None)?;
2050 overview.set_detail_availability(if has_worktree_aggregates {
2051 UsageDetailAvailability::Partial
2052 } else {
2053 native_detail
2054 });
2055 Ok(overview)
2056}
2057
2058pub(crate) fn worktree_token_overview(
2060 connection: &Connection,
2061 registration_id: i64,
2062) -> DbResult<TokenOverview> {
2063 let aggregates = load_worktree_overview_aggregates(connection, Some(registration_id), false)?;
2064 let (buckets, totals, average_policy_complete) = aggregate_report_rows(aggregates)?;
2065 let mut overview = TokenOverview::from_buckets(buckets);
2066 overview.apply_accounting_totals(totals);
2067 if !average_policy_complete {
2068 overview.average_policy.evidence = TOKEN_AVERAGE_POLICY_OVERFLOW_EVIDENCE.to_string();
2069 }
2070 overview.set_detail_availability(UsageDetailAvailability::Partial);
2071 Ok(overview)
2072}
2073
2074fn token_overview_for_project(
2076 connection: &Connection,
2077 project: ProjectInstanceId,
2078 caller_label: Option<&str>,
2079) -> DbResult<TokenOverview> {
2080 crate::project_identity::require_bound_project_identity(connection, project)?;
2081 let aggregates = load_overview_aggregates(connection, project, caller_label)?;
2082 let (buckets, totals, average_policy_complete) = aggregate_report_rows(aggregates)?;
2083 let mut overview = TokenOverview::from_buckets(buckets);
2084 overview.apply_accounting_totals(totals);
2085 if !average_policy_complete {
2086 overview.average_policy.evidence = TOKEN_AVERAGE_POLICY_OVERFLOW_EVIDENCE.to_string();
2087 }
2088 overview.set_detail_availability(detail_availability(connection, project, caller_label)?);
2089 Ok(overview)
2090}
2091
2092pub(crate) fn token_trends(
2094 connection: &Connection,
2095 caller_label: Option<&str>,
2096 window: TokenTrendWindow,
2097) -> DbResult<TokenTrendReport> {
2098 let project = current_project(connection)?;
2099 token_trends_for_project(connection, project, caller_label, window)
2100}
2101
2102pub(crate) fn repository_token_trends(
2104 connection: &Connection,
2105 window: TokenTrendWindow,
2106) -> DbResult<TokenTrendReport> {
2107 let project = current_project(connection)?;
2108 crate::project_identity::require_bound_project_identity(connection, project)?;
2109 let mut rows = load_daily_aggregates(connection, project, None, window)?;
2110 let worktree_rows = load_worktree_daily_aggregates(connection, None, true, window)?;
2111 let has_worktree_rows = worktree_aggregates_exist(connection)?;
2112 rows.extend(worktree_rows);
2113 token_trend_report(
2114 rows,
2115 None,
2116 window,
2117 if has_worktree_rows {
2118 UsageDetailAvailability::Partial
2119 } else {
2120 detail_availability(connection, project, None)?
2121 },
2122 )
2123}
2124
2125pub(crate) fn worktree_token_trends(
2127 connection: &Connection,
2128 registration_id: i64,
2129 window: TokenTrendWindow,
2130) -> DbResult<TokenTrendReport> {
2131 let rows = load_worktree_daily_aggregates(connection, Some(registration_id), false, window)?;
2132 token_trend_report(rows, None, window, UsageDetailAvailability::Partial)
2133}
2134
2135fn token_trends_for_project(
2137 connection: &Connection,
2138 project: ProjectInstanceId,
2139 caller_label: Option<&str>,
2140 window: TokenTrendWindow,
2141) -> DbResult<TokenTrendReport> {
2142 crate::project_identity::require_bound_project_identity(connection, project)?;
2143 let rows = load_daily_aggregates(connection, project, caller_label, window)?;
2144 token_trend_report(
2145 rows,
2146 caller_label.map(str::to_owned),
2147 window,
2148 detail_availability(connection, project, caller_label)?,
2149 )
2150}
2151
2152fn token_trend_report(
2154 rows: Vec<(String, DimensionValues, AggregateCounters)>,
2155 caller_label: Option<String>,
2156 window: TokenTrendWindow,
2157 detail: UsageDetailAvailability,
2158) -> DbResult<TokenTrendReport> {
2159 let mut by_period = BTreeMap::<String, BTreeMap<DimensionValues, AggregateCounters>>::new();
2160 for (period, dimension, counters) in rows {
2161 let entry = by_period
2162 .entry(period)
2163 .or_default()
2164 .entry(dimension)
2165 .or_default();
2166 *entry = entry.checked_add(counters)?;
2167 }
2168 let periods = by_period
2169 .into_iter()
2170 .map(|(period, rows)| {
2171 let buckets = rows
2172 .into_iter()
2173 .map(|(dimension, counters)| bucket_from_counters(dimension, counters))
2174 .collect::<DbResult<Vec<_>>>()?;
2175 Ok(TokenTrendPeriod::from_buckets(period, buckets))
2176 })
2177 .collect::<DbResult<Vec<_>>>()?;
2178 let mut report = TokenTrendReport::new(caller_label, window, periods);
2179 report.set_detail_availability(detail);
2180 Ok(report)
2181}
2182
2183fn current_project(connection: &Connection) -> DbResult<ProjectInstanceId> {
2185 crate::project_identity::load_project_identity(connection)?
2186 .ok_or(DbError::ProjectInstanceIdentityMissing)
2187}
2188
2189fn migrated_instance_id(
2191 project: ProjectInstanceId,
2192 caller_label: &str,
2193) -> DbResult<UsageInstanceId> {
2194 let mut hasher = blake3::Hasher::new();
2195 hasher.update(b"projectatlas:migrated-usage-instance:v2\0");
2196 hasher.update(&project.as_bytes());
2197 hasher.update(&(caller_label.len() as u64).to_le_bytes());
2198 hasher.update(caller_label.as_bytes());
2199 let mut bytes = [0_u8; 16];
2200 bytes.copy_from_slice(&hasher.finalize().as_bytes()[..16]);
2201 UsageInstanceId::from_bytes(bytes).map_err(Into::into)
2202}
2203
2204fn normalize_legacy_event(
2206 mut event: UsageEvent,
2207 policy: TelemetryRetentionPolicy,
2208) -> DbResult<(UsageEvent, LegacyDetailLoss)> {
2209 let baseline_identity = event.effective_baseline_identity().into_owned();
2210 let baseline_fingerprint = event.effective_baseline_fingerprint().into_owned();
2211 let mut loss = LegacyDetailLoss::default();
2212
2213 if !event.session_id.is_empty() {
2214 loss.label = normalize_legacy_required_text(
2215 "session_id",
2216 &mut event.session_id,
2217 policy.max_label_bytes,
2218 )?;
2219 }
2220 loss.raw |=
2221 normalize_legacy_required_text("command", &mut event.command, policy.max_command_bytes)?;
2222 loss.raw |= normalize_legacy_optional_text("path", &mut event.path, policy.max_path_bytes)?;
2223 loss.raw |= normalize_legacy_optional_text("query", &mut event.query, policy.max_query_bytes)?;
2224
2225 for (field, value) in [
2226 ("token_savings_bucket", &mut event.token_savings_bucket),
2227 ("provider", &mut event.provider),
2228 ("model", &mut event.model),
2229 ("tokenizer_backend", &mut event.tokenizer_backend),
2230 ("accuracy", &mut event.accuracy),
2231 ("baseline_kind", &mut event.baseline_kind),
2232 ("confidence", &mut event.confidence),
2233 ("accounting_layer", &mut event.accounting_layer),
2234 ("estimate_method", &mut event.estimate_method),
2235 ("denominator_kind", &mut event.denominator_kind),
2236 ("dedupe_scope", &mut event.dedupe_scope),
2237 ] {
2238 loss.dimension |= normalize_legacy_required_text(field, value, policy.max_dimension_bytes)?;
2239 }
2240
2241 loss.raw |= normalize_legacy_required_text(
2242 "calculation_trace",
2243 &mut event.calculation_trace,
2244 256.min(policy.max_baseline_witness_bytes),
2245 )?;
2246 event.baseline_identity = baseline_identity;
2247 loss.raw |= normalize_legacy_required_text(
2248 "baseline_identity",
2249 &mut event.baseline_identity,
2250 policy.max_baseline_witness_bytes,
2251 )?;
2252 event.baseline_fingerprint = baseline_fingerprint;
2253 loss.raw |= normalize_legacy_required_text(
2254 "baseline_fingerprint",
2255 &mut event.baseline_fingerprint,
2256 256.min(policy.max_baseline_witness_bytes),
2257 )?;
2258 Ok((event, loss))
2259}
2260
2261fn normalize_legacy_required_text(
2263 field: &'static str,
2264 value: &mut String,
2265 limit: usize,
2266) -> DbResult<bool> {
2267 if !value.is_empty() && value.len() <= limit {
2268 return Ok(false);
2269 }
2270 *value = legacy_text_token(field, value, limit)?;
2271 Ok(true)
2272}
2273
2274fn normalize_legacy_optional_text(
2276 field: &'static str,
2277 value: &mut Option<String>,
2278 limit: usize,
2279) -> DbResult<bool> {
2280 let Some(text) = value else {
2281 return Ok(false);
2282 };
2283 if text.len() <= limit {
2284 return Ok(false);
2285 }
2286 *text = legacy_text_token(field, text, limit)?;
2287 Ok(true)
2288}
2289
2290fn legacy_text_token(field: &'static str, value: &str, limit: usize) -> DbResult<String> {
2292 let mut hasher = blake3::Hasher::new();
2293 hasher.update(LEGACY_TEXT_HASH_DOMAIN);
2294 for bytes in [field.as_bytes(), value.as_bytes()] {
2295 hasher.update(&(bytes.len() as u64).to_le_bytes());
2296 hasher.update(bytes);
2297 }
2298 let token = format!("legacy:{field}:{}", hasher.finalize().to_hex());
2299 validate_required_text(field, &token, limit)?;
2300 Ok(token)
2301}
2302
2303fn mark_legacy_detail_loss(
2305 connection: &Connection,
2306 project: ProjectInstanceId,
2307 instance: UsageInstanceId,
2308 event: &UsageEvent,
2309 loss: LegacyDetailLoss,
2310) -> DbResult<()> {
2311 if loss == LegacyDetailLoss::default() {
2312 return Ok(());
2313 }
2314 if loss.raw {
2315 connection.execute(
2316 "UPDATE usage_instances SET raw_detail_complete = 0
2317 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2
2318 AND raw_detail_complete <> 0",
2319 params![
2320 project.as_bytes().as_slice(),
2321 instance.as_bytes().as_slice()
2322 ],
2323 )?;
2324 }
2325 if (loss.raw || loss.label)
2326 && let Some(label) = event_label(event)
2327 {
2328 connection.execute(
2329 "UPDATE usage_labels SET detail_complete = 0
2330 WHERE project_instance_id = ?1 AND caller_label = ?2
2331 AND detail_complete <> 0",
2332 params![project.as_bytes().as_slice(), label],
2333 )?;
2334 }
2335 connection.execute(
2336 "UPDATE usage_retention_state
2337 SET raw_detail_complete = CASE WHEN ?1 <> 0 THEN 0 ELSE raw_detail_complete END,
2338 dimension_detail_complete =
2339 CASE WHEN ?2 <> 0 THEN 0 ELSE dimension_detail_complete END,
2340 label_history_complete =
2341 CASE WHEN ?3 <> 0 THEN 0 ELSE label_history_complete END
2342 WHERE singleton = 1
2343 AND ((?1 <> 0 AND raw_detail_complete <> 0)
2344 OR (?2 <> 0 AND dimension_detail_complete <> 0)
2345 OR (?3 <> 0 AND label_history_complete <> 0))",
2346 params![
2347 i64::from(loss.raw),
2348 i64::from(loss.dimension),
2349 i64::from(loss.label),
2350 ],
2351 )?;
2352 Ok(())
2353}
2354
2355fn event_label(event: &UsageEvent) -> Option<&str> {
2357 (!event.session_id.is_empty()).then_some(event.session_id.as_str())
2358}
2359
2360pub(crate) fn validate_event(event: &UsageEvent, policy: TelemetryRetentionPolicy) -> DbResult<()> {
2362 validate_optional_text("session_id", event_label(event), policy.max_label_bytes)?;
2363 validate_required_text("command", &event.command, policy.max_command_bytes)?;
2364 validate_optional_text("path", event.path.as_deref(), policy.max_path_bytes)?;
2365 validate_optional_text("query", event.query.as_deref(), policy.max_query_bytes)?;
2366 for (field, value) in [
2367 ("token_savings_bucket", event.token_savings_bucket.as_str()),
2368 ("provider", event.provider.as_str()),
2369 ("model", event.model.as_str()),
2370 ("tokenizer_backend", event.tokenizer_backend.as_str()),
2371 ("accuracy", event.accuracy.as_str()),
2372 ("baseline_kind", event.baseline_kind.as_str()),
2373 ("confidence", event.confidence.as_str()),
2374 ("accounting_layer", event.report_accounting_layer()),
2375 ("estimate_method", event.estimate_method.as_str()),
2376 ("denominator_kind", event.report_denominator_kind()),
2377 ("dedupe_scope", event.report_dedupe_scope()),
2378 ] {
2379 validate_required_text(field, value, policy.max_dimension_bytes)?;
2380 }
2381 validate_required_text(
2382 "calculation_trace",
2383 &event.calculation_trace,
2384 256.min(policy.max_baseline_witness_bytes),
2385 )?;
2386 let baseline_identity = event.effective_baseline_identity();
2387 let baseline_fingerprint = event.effective_baseline_fingerprint();
2388 validate_required_text(
2389 "baseline_identity",
2390 baseline_identity.as_ref(),
2391 policy.max_baseline_witness_bytes,
2392 )?;
2393 validate_required_text(
2394 "baseline_fingerprint",
2395 baseline_fingerprint.as_ref(),
2396 256.min(policy.max_baseline_witness_bytes),
2397 )?;
2398 let _ = option_usize_to_i64(
2399 "estimated_tokens_without_projectatlas",
2400 event.estimated_tokens_without_projectatlas,
2401 )?;
2402 let _ = option_usize_to_i64(
2403 "estimated_tokens_with_projectatlas",
2404 event.estimated_tokens_with_projectatlas,
2405 )?;
2406 let _ = option_isize_to_i64("estimated_tokens_saved", event.estimated_tokens_saved)?;
2407 Ok(())
2408}
2409
2410fn validate_required_text(field: &'static str, value: &str, limit: usize) -> DbResult<()> {
2412 if value.is_empty() || value.len() > limit {
2413 return Err(DbError::TelemetryFieldTooLarge {
2414 field,
2415 bytes: value.len(),
2416 limit,
2417 });
2418 }
2419 Ok(())
2420}
2421
2422fn validate_optional_text(field: &'static str, value: Option<&str>, limit: usize) -> DbResult<()> {
2424 if let Some(value) = value
2425 && value.len() > limit
2426 {
2427 return Err(DbError::TelemetryFieldTooLarge {
2428 field,
2429 bytes: value.len(),
2430 limit,
2431 });
2432 }
2433 Ok(())
2434}
2435
2436fn ensure_overflow_dimension(connection: &Connection, denominator_kind: &str) -> DbResult<i64> {
2438 ensure_dimension_unbounded(connection, &DimensionValues::overflow(denominator_kind))
2439}
2440
2441fn ensure_dimension(
2443 connection: &Connection,
2444 dimension: &DimensionValues,
2445 policy: TelemetryRetentionPolicy,
2446) -> DbResult<i64> {
2447 if let Some(id) = find_dimension(connection, dimension)? {
2448 return Ok(id);
2449 }
2450 if retention_counter(connection, RetentionCounter::DimensionRows)? >= policy.max_dimensions {
2451 connection.execute(
2452 "UPDATE usage_retention_state SET dimension_detail_complete = 0 WHERE singleton = 1",
2453 [],
2454 )?;
2455 return ensure_overflow_dimension(connection, &dimension.denominator_kind);
2456 }
2457 ensure_dimension_unbounded(connection, dimension)
2458}
2459
2460fn ensure_dimension_unbounded(
2462 connection: &Connection,
2463 dimension: &DimensionValues,
2464) -> DbResult<i64> {
2465 let inserted = connection.execute(
2466 "INSERT OR IGNORE INTO usage_bucket_dimensions(
2467 token_savings_bucket, provider, model, tokenizer_backend,
2468 accuracy, baseline_kind, confidence, accounting_layer,
2469 estimate_method, denominator_kind, dedupe_scope, overflow
2470 ) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)",
2471 params![
2472 dimension.token_savings_bucket,
2473 dimension.provider,
2474 dimension.model,
2475 dimension.tokenizer_backend,
2476 dimension.accuracy,
2477 dimension.baseline_kind,
2478 dimension.confidence,
2479 dimension.accounting_layer,
2480 dimension.estimate_method,
2481 dimension.denominator_kind,
2482 dimension.dedupe_scope,
2483 i64::from(dimension.overflow),
2484 ],
2485 )?;
2486 if inserted != 0 {
2487 increment_retention_counter(connection, RetentionCounter::DimensionRows, 1)?;
2488 }
2489 find_dimension(connection, dimension)?.ok_or(DbError::SchemaPostcondition { expected: 11 })
2490}
2491
2492fn find_dimension(connection: &Connection, dimension: &DimensionValues) -> DbResult<Option<i64>> {
2494 connection
2495 .query_row(
2496 "SELECT dimension_id FROM usage_bucket_dimensions
2497 WHERE token_savings_bucket = ?1 AND provider = ?2 AND model = ?3
2498 AND tokenizer_backend = ?4 AND accuracy = ?5
2499 AND baseline_kind = ?6 AND confidence = ?7
2500 AND accounting_layer = ?8 AND estimate_method = ?9
2501 AND denominator_kind = ?10 AND dedupe_scope = ?11
2502 AND overflow = ?12",
2503 params![
2504 dimension.token_savings_bucket,
2505 dimension.provider,
2506 dimension.model,
2507 dimension.tokenizer_backend,
2508 dimension.accuracy,
2509 dimension.baseline_kind,
2510 dimension.confidence,
2511 dimension.accounting_layer,
2512 dimension.estimate_method,
2513 dimension.denominator_kind,
2514 dimension.dedupe_scope,
2515 i64::from(dimension.overflow),
2516 ],
2517 |row| row.get(0),
2518 )
2519 .optional()
2520 .map_err(Into::into)
2521}
2522
2523fn ensure_active_instance(
2525 connection: &Connection,
2526 project: ProjectInstanceId,
2527 runtime: UsageInstanceId,
2528 owner: UsageInstanceOwner,
2529 caller_label: Option<&str>,
2530 policy: TelemetryRetentionPolicy,
2531 now: i64,
2532) -> DbResult<i64> {
2533 let existing = connection
2534 .query_row(
2535 "SELECT instance_row_id, owner, caller_label, state
2536 FROM usage_instances
2537 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2",
2538 params![project.as_bytes().as_slice(), runtime.as_bytes().as_slice()],
2539 |row| {
2540 Ok((
2541 row.get::<_, i64>(0)?,
2542 row.get::<_, String>(1)?,
2543 row.get::<_, Option<String>>(2)?,
2544 row.get::<_, String>(3)?,
2545 ))
2546 },
2547 )
2548 .optional()?;
2549 if let Some((row_id, stored_owner, stored_label, state)) = existing {
2550 if stored_owner != owner.as_str() || stored_label.as_deref() != caller_label {
2551 return Err(DbError::TelemetryInstanceMismatch);
2552 }
2553 if state != INSTANCE_ACTIVE {
2554 return Err(DbError::TelemetryInstanceInactive);
2555 }
2556 return Ok(row_id);
2557 }
2558 let retired = connection.query_row(
2559 "SELECT EXISTS(
2560 SELECT 1 FROM usage_instance_tombstones
2561 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2
2562 )",
2563 params![project.as_bytes().as_slice(), runtime.as_bytes().as_slice()],
2564 |row| row.get::<_, i64>(0),
2565 )?;
2566 if retired != 0 {
2567 return Err(DbError::TelemetryInstanceInactive);
2568 }
2569 let active = connection.query_row(
2570 "SELECT COUNT(*) FROM usage_instances
2571 WHERE project_instance_id = ?1 AND state = ?2",
2572 params![project.as_bytes().as_slice(), INSTANCE_ACTIVE],
2573 |row| row.get::<_, i64>(0),
2574 )?;
2575 if count_usize("active_usage_instances", active)? >= policy.max_active_instances {
2576 return Err(DbError::TelemetryInstanceCapacity);
2577 }
2578 if retention_counter(connection, RetentionCounter::InstanceRows)?
2579 >= policy.max_retained_instances
2580 {
2581 return Err(DbError::TelemetryInstanceCapacity);
2582 }
2583 connection.execute(
2584 "INSERT INTO usage_instances(
2585 project_instance_id, runtime_instance_id, owner, caller_label, state,
2586 started_at_epoch, last_seen_at_epoch
2587 ) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?6)",
2588 params![
2589 project.as_bytes().as_slice(),
2590 runtime.as_bytes().as_slice(),
2591 owner.as_str(),
2592 caller_label,
2593 INSTANCE_ACTIVE,
2594 now,
2595 ],
2596 )?;
2597 increment_retention_counter(connection, RetentionCounter::InstanceRows, 1)?;
2598 Ok(connection.last_insert_rowid())
2599}
2600
2601fn ensure_label(
2603 connection: &Connection,
2604 project: ProjectInstanceId,
2605 caller_label: Option<&str>,
2606 policy: TelemetryRetentionPolicy,
2607 now: i64,
2608) -> DbResult<()> {
2609 let Some(label) = caller_label else {
2610 return Ok(());
2611 };
2612 let exists = connection.query_row(
2613 "SELECT EXISTS(
2614 SELECT 1 FROM usage_labels
2615 WHERE project_instance_id = ?1 AND caller_label = ?2
2616 )",
2617 params![project.as_bytes().as_slice(), label],
2618 |row| row.get::<_, i64>(0),
2619 )?;
2620 let prior_history = connection.query_row(
2621 "SELECT EXISTS(
2622 SELECT 1 FROM usage_label_tombstones
2623 WHERE project_instance_id = ?1 AND caller_label = ?2
2624 )",
2625 params![project.as_bytes().as_slice(), label],
2626 |row| row.get::<_, i64>(0),
2627 )?;
2628 if exists == 0
2629 && retention_counter(connection, RetentionCounter::LabelRows)? >= policy.max_retained_labels
2630 && !evict_oldest_label(connection, policy, now)?
2631 {
2632 return Err(DbError::TelemetryInstanceCapacity);
2633 }
2634 let inserted = connection.execute(
2635 "INSERT INTO usage_labels(
2636 project_instance_id, caller_label, last_seen_at_epoch, detail_complete
2637 ) VALUES(?1, ?2, ?3, ?4)
2638 ON CONFLICT(project_instance_id, caller_label) DO UPDATE SET
2639 last_seen_at_epoch = excluded.last_seen_at_epoch",
2640 params![
2641 project.as_bytes().as_slice(),
2642 label,
2643 now,
2644 i64::from(prior_history == 0),
2645 ],
2646 )?;
2647 if inserted != 0 && exists == 0 {
2648 increment_retention_counter(connection, RetentionCounter::LabelRows, 1)?;
2649 }
2650 Ok(())
2651}
2652
2653fn upsert_label_tombstone(
2655 connection: &Connection,
2656 project: &[u8],
2657 label: &str,
2658 now: i64,
2659 runtime: Option<&[u8]>,
2660) -> DbResult<()> {
2661 let exists = connection.query_row(
2662 "SELECT EXISTS(
2663 SELECT 1 FROM usage_label_tombstones
2664 WHERE project_instance_id = ?1 AND caller_label = ?2
2665 )",
2666 params![project, label],
2667 |row| row.get::<_, i64>(0),
2668 )?;
2669 connection.execute(
2670 "INSERT INTO usage_label_tombstones(
2671 project_instance_id, caller_label, expired_at_epoch, last_instance_id
2672 ) VALUES(?1, ?2, ?3, ?4)
2673 ON CONFLICT(project_instance_id, caller_label) DO UPDATE SET
2674 expired_at_epoch = excluded.expired_at_epoch,
2675 last_instance_id = excluded.last_instance_id",
2676 params![project, label, now, runtime],
2677 )?;
2678 if exists == 0 {
2679 increment_retention_counter(connection, RetentionCounter::LabelTombstoneRows, 1)?;
2680 }
2681 Ok(())
2682}
2683
2684fn evict_oldest_label(
2686 connection: &Connection,
2687 policy: TelemetryRetentionPolicy,
2688 now: i64,
2689) -> DbResult<bool> {
2690 let candidate = connection
2691 .query_row(
2692 "SELECT project_instance_id, caller_label FROM usage_labels
2693 WHERE NOT EXISTS(
2694 SELECT 1 FROM usage_instances AS i
2695 WHERE i.project_instance_id = usage_labels.project_instance_id
2696 AND i.caller_label = usage_labels.caller_label
2697 AND i.state = ?1
2698 )
2699 AND (
2700 SELECT COUNT(*) FROM usage_instances AS i
2701 WHERE i.project_instance_id = usage_labels.project_instance_id
2702 AND i.caller_label = usage_labels.caller_label
2703 ) <= ?2
2704 ORDER BY last_seen_at_epoch, project_instance_id, caller_label LIMIT 1",
2705 params![
2706 INSTANCE_ACTIVE,
2707 to_i64("prune_batch_rows", policy.prune_batch_rows)?,
2708 ],
2709 |row| Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, String>(1)?)),
2710 )
2711 .optional()?;
2712 let Some((project, label)) = candidate else {
2713 return Ok(false);
2714 };
2715 let runtime = connection
2716 .query_row(
2717 "SELECT runtime_instance_id FROM usage_instances
2718 WHERE project_instance_id = ?1 AND caller_label = ?2
2719 ORDER BY last_seen_at_epoch DESC, instance_row_id DESC LIMIT 1",
2720 params![project, label],
2721 |row| row.get::<_, Vec<u8>>(0),
2722 )
2723 .optional()?;
2724 upsert_label_tombstone(
2725 connection,
2726 project.as_slice(),
2727 &label,
2728 now,
2729 runtime.as_deref(),
2730 )?;
2731 connection.execute(
2732 "UPDATE usage_instances SET caller_label = NULL
2733 WHERE project_instance_id = ?1 AND caller_label = ?2 AND state <> ?3",
2734 params![project, label, INSTANCE_ACTIVE],
2735 )?;
2736 let deleted = connection.execute(
2737 "DELETE FROM usage_labels
2738 WHERE project_instance_id = ?1 AND caller_label = ?2",
2739 params![project, label],
2740 )?;
2741 decrement_retention_counter(connection, RetentionCounter::LabelRows, deleted)?;
2742 connection.execute(
2743 "UPDATE usage_retention_state SET label_history_complete = 0 WHERE singleton = 1",
2744 [],
2745 )?;
2746 Ok(true)
2747}
2748
2749fn aggregate_delta(
2751 connection: &Connection,
2752 instance_row_id: i64,
2753 event: &UsageEvent,
2754 policy: TelemetryRetentionPolicy,
2755 baseline_admission: BaselineAdmission,
2756) -> DbResult<AggregateCounters> {
2757 let (Some(without_source), Some(with_source)) = (
2758 event.estimated_tokens_without_projectatlas,
2759 event.estimated_tokens_with_projectatlas,
2760 ) else {
2761 return Ok(AggregateCounters::default());
2762 };
2763 let without = to_i64("estimated_tokens_without_projectatlas", without_source)?;
2764 let with = to_i64("estimated_tokens_with_projectatlas", with_source)?;
2765 let observed = event.is_observed();
2766 let modeled = event.is_modeled();
2767 let mut counters = AggregateCounters {
2768 calls: 1,
2769 estimated_without: without,
2770 estimated_with: with,
2771 observed_without: if observed { without } else { 0 },
2772 observed_with: if observed { with } else { 0 },
2773 modeled_without: if modeled { without } else { 0 },
2774 modeled_with: if modeled { with } else { 0 },
2775 observed_file_read_replacements: i64::from(
2776 observed && event.is_observed_file_read_replacement(without_source),
2777 ),
2778 modeled_file_reads_avoided: i64::from(
2779 modeled && event.is_modeled_file_read_avoidance(without_source),
2780 ),
2781 ..AggregateCounters::default()
2782 };
2783 if modeled {
2784 let contribution = if event.report_dedupe_scope() == DEDUPE_SCOPE_EVENT {
2785 without
2786 .checked_sub(with)
2787 .ok_or(DbError::TelemetryIntegerOverflow {
2788 field: "event_modeled_contribution",
2789 })?
2790 } else {
2791 let (adjustment, repeated) = update_modeled_baseline(
2792 connection,
2793 instance_row_id,
2794 event,
2795 without,
2796 with,
2797 policy,
2798 baseline_admission,
2799 )?;
2800 counters.repeated_baselines = repeated;
2801 adjustment
2802 };
2803 let (positive, negative) = signed_components(contribution)?;
2804 counters.deduped_modeled_without = positive;
2805 counters.deduped_modeled_with = negative;
2806 }
2807 Ok(counters)
2808}
2809
2810fn update_modeled_baseline(
2812 connection: &Connection,
2813 instance_row_id: i64,
2814 event: &UsageEvent,
2815 without: i64,
2816 with: i64,
2817 policy: TelemetryRetentionPolicy,
2818 baseline_admission: BaselineAdmission,
2819) -> DbResult<(i64, i64)> {
2820 let key = event.modeled_baseline_key();
2821 let identity = event.effective_baseline_identity();
2822 let fingerprint = event.effective_baseline_fingerprint();
2823 let existing = connection
2824 .query_row(
2825 "SELECT baseline_identity, baseline_fingerprint, denominator_kind,
2826 maximum_without, emitted_with, calls
2827 FROM usage_instance_baselines
2828 WHERE instance_row_id = ?1 AND baseline_key = ?2",
2829 params![instance_row_id, key.as_slice()],
2830 |row| {
2831 Ok((
2832 row.get::<_, String>(0)?,
2833 row.get::<_, String>(1)?,
2834 row.get::<_, String>(2)?,
2835 row.get::<_, i64>(3)?,
2836 row.get::<_, i64>(4)?,
2837 row.get::<_, i64>(5)?,
2838 ))
2839 },
2840 )
2841 .optional()?;
2842 if let Some((stored_identity, stored_fingerprint, denominator, old_max, old_with, calls)) =
2843 existing
2844 {
2845 if stored_identity != identity
2846 || stored_fingerprint != fingerprint
2847 || denominator != event.report_denominator_kind()
2848 {
2849 return Err(DbError::TelemetryBaselineCollision);
2850 }
2851 let previous = old_max
2852 .checked_sub(old_with)
2853 .ok_or(DbError::TelemetryIntegerOverflow {
2854 field: "previous_modeled_baseline",
2855 })?;
2856 let maximum = old_max.max(without);
2857 let emitted = old_with
2858 .checked_add(with)
2859 .ok_or(DbError::TelemetryIntegerOverflow {
2860 field: "modeled_baseline_emitted_with",
2861 })?;
2862 let current = maximum
2863 .checked_sub(emitted)
2864 .ok_or(DbError::TelemetryIntegerOverflow {
2865 field: "current_modeled_baseline",
2866 })?;
2867 let adjustment =
2868 current
2869 .checked_sub(previous)
2870 .ok_or(DbError::TelemetryIntegerOverflow {
2871 field: "modeled_baseline_adjustment",
2872 })?;
2873 let calls = calls
2874 .checked_add(1)
2875 .ok_or(DbError::TelemetryIntegerOverflow {
2876 field: "modeled_baseline_calls",
2877 })?;
2878 connection.execute(
2879 "UPDATE usage_instance_baselines
2880 SET maximum_without = ?3, emitted_with = ?4, calls = ?5
2881 WHERE instance_row_id = ?1 AND baseline_key = ?2",
2882 params![instance_row_id, key.as_slice(), maximum, emitted, calls],
2883 )?;
2884 return Ok((adjustment, 1));
2885 }
2886 let instance_count = connection.query_row(
2887 "SELECT COUNT(*) FROM usage_instance_baselines WHERE instance_row_id = ?1",
2888 [instance_row_id],
2889 |row| row.get::<_, i64>(0),
2890 )?;
2891 let total_count = retention_counter(connection, RetentionCounter::BaselineRows)?;
2892 let total_bytes = retention_counter(connection, RetentionCounter::BaselineLogicalBytes)?;
2893 let witness_bytes = identity
2894 .len()
2895 .checked_add(fingerprint.len())
2896 .and_then(|value| value.checked_add(event.report_denominator_kind().len()))
2897 .ok_or(DbError::TelemetryIntegerOverflow {
2898 field: "baseline_witness_logical_bytes",
2899 })?;
2900 let projected_bytes =
2901 total_bytes
2902 .checked_add(witness_bytes)
2903 .ok_or(DbError::TelemetryIntegerOverflow {
2904 field: "baseline_logical_bytes",
2905 })?;
2906 let exceeds_runtime_limit = count_usize("instance_baseline_rows", instance_count)?
2907 >= policy.max_baselines_per_instance
2908 || total_count >= policy.max_active_baseline_rows
2909 || projected_bytes > policy.max_baseline_logical_bytes;
2910 if baseline_admission == BaselineAdmission::BoundedRuntime && exceeds_runtime_limit {
2911 return Err(DbError::TelemetryBaselineCapacity);
2912 }
2913 connection.execute(
2914 "INSERT INTO usage_instance_baselines(
2915 instance_row_id, baseline_key, baseline_identity, baseline_fingerprint,
2916 denominator_kind, maximum_without, emitted_with, calls, witness_logical_bytes
2917 ) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, 1, ?8)",
2918 params![
2919 instance_row_id,
2920 key.as_slice(),
2921 identity,
2922 fingerprint,
2923 event.report_denominator_kind(),
2924 without,
2925 with,
2926 to_i64("baseline_witness_logical_bytes", witness_bytes)?,
2927 ],
2928 )?;
2929 increment_retention_counter(connection, RetentionCounter::BaselineRows, 1)?;
2930 increment_retention_counter(
2931 connection,
2932 RetentionCounter::BaselineLogicalBytes,
2933 witness_bytes,
2934 )?;
2935 Ok((
2936 without
2937 .checked_sub(with)
2938 .ok_or(DbError::TelemetryIntegerOverflow {
2939 field: "initial_modeled_baseline",
2940 })?,
2941 0,
2942 ))
2943}
2944
2945fn delete_instance_baselines(connection: &Connection, instance_row_id: i64) -> DbResult<usize> {
2947 let (rows, bytes) = connection.query_row(
2948 "SELECT COUNT(*), COALESCE(SUM(witness_logical_bytes), 0)
2949 FROM usage_instance_baselines WHERE instance_row_id = ?1",
2950 [instance_row_id],
2951 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
2952 )?;
2953 let rows = count_usize("instance_baseline_rows", rows)?;
2954 let bytes = count_usize("instance_baseline_logical_bytes", bytes)?;
2955 if rows != 0 {
2956 connection.execute(
2957 "DELETE FROM usage_instance_baselines WHERE instance_row_id = ?1",
2958 [instance_row_id],
2959 )?;
2960 decrement_retention_counter(connection, RetentionCounter::BaselineRows, rows)?;
2961 decrement_retention_counter(connection, RetentionCounter::BaselineLogicalBytes, bytes)?;
2962 }
2963 Ok(rows)
2964}
2965
2966fn insert_raw_event(
2968 connection: &Connection,
2969 instance_row_id: i64,
2970 dimension_id: i64,
2971 event: &UsageEvent,
2972 created_at: i64,
2973 logical_bytes: usize,
2974) -> DbResult<()> {
2975 connection.execute(
2976 "INSERT INTO usage_events(
2977 instance_row_id, dimension_id, command, path, query,
2978 estimated_tokens_without_projectatlas,
2979 estimated_tokens_with_projectatlas, estimated_tokens_saved,
2980 calculation_trace, baseline_identity, baseline_fingerprint,
2981 created_at_epoch, logical_bytes
2982 ) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)",
2983 params![
2984 instance_row_id,
2985 dimension_id,
2986 event.command,
2987 event.path,
2988 event.query,
2989 option_usize_to_i64(
2990 "estimated_tokens_without_projectatlas",
2991 event.estimated_tokens_without_projectatlas,
2992 )?,
2993 option_usize_to_i64(
2994 "estimated_tokens_with_projectatlas",
2995 event.estimated_tokens_with_projectatlas,
2996 )?,
2997 option_isize_to_i64("estimated_tokens_saved", event.estimated_tokens_saved)?,
2998 event.calculation_trace,
2999 event.effective_baseline_identity(),
3000 event.effective_baseline_fingerprint(),
3001 created_at,
3002 to_i64("raw_logical_bytes", logical_bytes)?,
3003 ],
3004 )?;
3005 increment_retention_counter(connection, RetentionCounter::RawRows, 1)?;
3006 increment_retention_counter(connection, RetentionCounter::RawLogicalBytes, logical_bytes)?;
3007 Ok(())
3008}
3009
3010fn bind_worktree_origin(
3012 connection: &Connection,
3013 instance_row_id: i64,
3014 registration_id: Option<i64>,
3015) -> DbResult<()> {
3016 let existing = connection
3017 .query_row(
3018 "SELECT registration_id FROM usage_instance_worktree_origins
3019 WHERE instance_row_id = ?1",
3020 [instance_row_id],
3021 |row| row.get::<_, i64>(0),
3022 )
3023 .optional()?;
3024 let Some(registration_id) = registration_id else {
3025 return if existing.is_some() {
3026 Err(DbError::WorktreeTelemetryOriginConflict)
3027 } else {
3028 Ok(())
3029 };
3030 };
3031 let registration_exists = connection.query_row(
3032 "SELECT EXISTS(
3033 SELECT 1 FROM worktree_registrations WHERE registration_id = ?1
3034 )",
3035 [registration_id],
3036 |row| row.get::<_, bool>(0),
3037 )?;
3038 if !registration_exists {
3039 return Err(DbError::WorktreeTelemetryProjectMismatch { registration_id });
3040 }
3041 match existing {
3042 Some(existing_id) if existing_id != registration_id => {
3043 Err(DbError::WorktreeTelemetryOriginConflict)
3044 }
3045 Some(_) => Ok(()),
3046 None => {
3047 connection.execute(
3048 "INSERT INTO usage_instance_worktree_origins(instance_row_id, registration_id)
3049 VALUES(?1, ?2)",
3050 params![instance_row_id, registration_id],
3051 )?;
3052 Ok(())
3053 }
3054 }
3055}
3056
3057fn increment_aggregate_revision(
3059 connection: &Connection,
3060 project: ProjectInstanceId,
3061 calls: i64,
3062) -> DbResult<()> {
3063 if calls == 0 {
3064 return Ok(());
3065 }
3066 connection
3067 .execute(
3068 "INSERT INTO usage_aggregate_revisions(project_instance_id, revision)
3069 VALUES(?1, 1)
3070 ON CONFLICT(project_instance_id) DO UPDATE SET
3071 revision = usage_aggregate_revisions.revision + 1",
3072 [project.as_bytes().as_slice()],
3073 )
3074 .map_err(aggregate_write_error)?;
3075 Ok(())
3076}
3077
3078fn upsert_routed_worktree_aggregates(
3080 connection: &Connection,
3081 registration_id: i64,
3082 dimension_id: i64,
3083 created_at: i64,
3084 delta: AggregateCounters,
3085 policy: TelemetryRetentionPolicy,
3086) -> DbResult<()> {
3087 if delta.calls == 0 {
3088 return Ok(());
3089 }
3090 upsert_worktree_aggregate(
3091 connection,
3092 registration_id,
3093 WORKTREE_USAGE_ROUTED,
3094 -1,
3095 dimension_id,
3096 delta,
3097 )?;
3098 let day = created_at - created_at.rem_euclid(SECONDS_PER_DAY);
3099 reserve_routed_worktree_daily_row(
3100 connection,
3101 registration_id,
3102 day,
3103 dimension_id,
3104 policy.max_daily_rows,
3105 )?;
3106 upsert_worktree_aggregate(
3107 connection,
3108 registration_id,
3109 WORKTREE_USAGE_ROUTED,
3110 day,
3111 dimension_id,
3112 delta,
3113 )
3114}
3115
3116fn upsert_worktree_aggregate(
3118 connection: &Connection,
3119 registration_id: i64,
3120 source_kind: &str,
3121 day_epoch: i64,
3122 dimension_id: i64,
3123 delta: AggregateCounters,
3124) -> DbResult<()> {
3125 connection
3126 .execute(
3127 "INSERT INTO worktree_usage_aggregates(
3128 registration_id, source_kind, day_epoch, dimension_id,
3129 calls, estimated_without, estimated_with, observed_without,
3130 observed_with, modeled_without, modeled_with,
3131 deduped_modeled_without, deduped_modeled_with, repeated_baselines,
3132 observed_file_read_replacements, modeled_file_reads_avoided
3133 ) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)
3134 ON CONFLICT(registration_id, source_kind, day_epoch, dimension_id) DO UPDATE SET
3135 calls = worktree_usage_aggregates.calls + excluded.calls,
3136 estimated_without = worktree_usage_aggregates.estimated_without + excluded.estimated_without,
3137 estimated_with = worktree_usage_aggregates.estimated_with + excluded.estimated_with,
3138 observed_without = worktree_usage_aggregates.observed_without + excluded.observed_without,
3139 observed_with = worktree_usage_aggregates.observed_with + excluded.observed_with,
3140 modeled_without = worktree_usage_aggregates.modeled_without + excluded.modeled_without,
3141 modeled_with = worktree_usage_aggregates.modeled_with + excluded.modeled_with,
3142 deduped_modeled_without = worktree_usage_aggregates.deduped_modeled_without + excluded.deduped_modeled_without,
3143 deduped_modeled_with = worktree_usage_aggregates.deduped_modeled_with + excluded.deduped_modeled_with,
3144 repeated_baselines = worktree_usage_aggregates.repeated_baselines + excluded.repeated_baselines,
3145 observed_file_read_replacements = worktree_usage_aggregates.observed_file_read_replacements + excluded.observed_file_read_replacements,
3146 modeled_file_reads_avoided = worktree_usage_aggregates.modeled_file_reads_avoided + excluded.modeled_file_reads_avoided",
3147 worktree_aggregate_params!(
3148 registration_id,
3149 source_kind,
3150 day_epoch,
3151 dimension_id,
3152 delta
3153 ),
3154 )
3155 .map_err(aggregate_write_error)?;
3156 Ok(())
3157}
3158
3159fn reserve_routed_worktree_daily_row(
3161 connection: &Connection,
3162 registration_id: i64,
3163 day_epoch: i64,
3164 dimension_id: i64,
3165 max_daily_rows: usize,
3166) -> DbResult<()> {
3167 let exists = connection.query_row(
3168 "SELECT EXISTS(
3169 SELECT 1 FROM worktree_usage_aggregates
3170 WHERE registration_id = ?1 AND source_kind = ?2
3171 AND day_epoch = ?3 AND dimension_id = ?4
3172 )",
3173 params![
3174 registration_id,
3175 WORKTREE_USAGE_ROUTED,
3176 day_epoch,
3177 dimension_id
3178 ],
3179 |row| row.get::<_, bool>(0),
3180 )?;
3181 if exists {
3182 return Ok(());
3183 }
3184 let daily_rows = connection.query_row(
3185 "SELECT COUNT(*) FROM worktree_usage_aggregates WHERE day_epoch >= 0",
3186 [],
3187 |row| row.get::<_, i64>(0),
3188 )?;
3189 let daily_rows = count_usize("worktree_daily_rows", daily_rows)?;
3190 if daily_rows < max_daily_rows {
3191 return Ok(());
3192 }
3193 let removed = connection.execute(
3194 "DELETE FROM worktree_usage_aggregates
3195 WHERE (registration_id, source_kind, day_epoch, dimension_id) IN (
3196 SELECT registration_id, source_kind, day_epoch, dimension_id
3197 FROM worktree_usage_aggregates
3198 WHERE day_epoch >= 0 AND source_kind = ?2
3199 AND NOT (
3200 registration_id = ?1 AND source_kind = ?2
3201 AND day_epoch = ?3 AND dimension_id = ?4
3202 )
3203 ORDER BY day_epoch, registration_id, source_kind, dimension_id
3204 LIMIT 1
3205 )",
3206 params![
3207 registration_id,
3208 WORKTREE_USAGE_ROUTED,
3209 day_epoch,
3210 dimension_id
3211 ],
3212 )?;
3213 if removed != 1 {
3214 return Err(DbError::WorktreeTelemetrySnapshotLimit {
3215 resource: "daily_rows",
3216 limit: max_daily_rows,
3217 observed: daily_rows.saturating_add(1),
3218 });
3219 }
3220 Ok(())
3221}
3222
3223fn apply_aggregates(
3225 connection: &Connection,
3226 project: ProjectInstanceId,
3227 instance_row_id: i64,
3228 dimension_id: i64,
3229 created_at: i64,
3230 delta: AggregateCounters,
3231 policy: TelemetryRetentionPolicy,
3232) -> DbResult<()> {
3233 if delta.calls == 0 {
3234 return Ok(());
3235 }
3236 upsert_global(connection, project, dimension_id, delta)?;
3237 upsert_instance(connection, instance_row_id, dimension_id, delta)?;
3238 let day = created_at - created_at.rem_euclid(SECONDS_PER_DAY);
3239 prepare_daily_capacity(
3240 connection,
3241 project,
3242 instance_row_id,
3243 day,
3244 dimension_id,
3245 policy,
3246 )?;
3247 upsert_global_daily(connection, project, day, dimension_id, delta)?;
3248 upsert_instance_daily(connection, instance_row_id, day, dimension_id, delta)?;
3249 Ok(())
3250}
3251
3252fn upsert_global(
3254 connection: &Connection,
3255 project: ProjectInstanceId,
3256 dimension_id: i64,
3257 delta: AggregateCounters,
3258) -> DbResult<()> {
3259 connection.execute(
3260 "INSERT INTO usage_global_aggregates VALUES(
3261 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
3262 ON CONFLICT(project_instance_id, dimension_id) DO UPDATE SET
3263 calls = usage_global_aggregates.calls + excluded.calls,
3264 estimated_without = usage_global_aggregates.estimated_without + excluded.estimated_without,
3265 estimated_with = usage_global_aggregates.estimated_with + excluded.estimated_with,
3266 observed_without = usage_global_aggregates.observed_without + excluded.observed_without,
3267 observed_with = usage_global_aggregates.observed_with + excluded.observed_with,
3268 modeled_without = usage_global_aggregates.modeled_without + excluded.modeled_without,
3269 modeled_with = usage_global_aggregates.modeled_with + excluded.modeled_with,
3270 deduped_modeled_without = usage_global_aggregates.deduped_modeled_without + excluded.deduped_modeled_without,
3271 deduped_modeled_with = usage_global_aggregates.deduped_modeled_with + excluded.deduped_modeled_with,
3272 repeated_baselines = usage_global_aggregates.repeated_baselines + excluded.repeated_baselines,
3273 observed_file_read_replacements = usage_global_aggregates.observed_file_read_replacements + excluded.observed_file_read_replacements,
3274 modeled_file_reads_avoided = usage_global_aggregates.modeled_file_reads_avoided + excluded.modeled_file_reads_avoided",
3275 aggregate_params!(project.as_bytes().as_slice(), dimension_id, delta),
3276 )
3277 .map_err(aggregate_write_error)?;
3278 Ok(())
3279}
3280
3281fn upsert_instance(
3283 connection: &Connection,
3284 instance_row_id: i64,
3285 dimension_id: i64,
3286 delta: AggregateCounters,
3287) -> DbResult<()> {
3288 connection.execute(
3289 "INSERT INTO usage_instance_aggregates VALUES(
3290 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
3291 ON CONFLICT(instance_row_id, dimension_id) DO UPDATE SET
3292 calls = usage_instance_aggregates.calls + excluded.calls,
3293 estimated_without = usage_instance_aggregates.estimated_without + excluded.estimated_without,
3294 estimated_with = usage_instance_aggregates.estimated_with + excluded.estimated_with,
3295 observed_without = usage_instance_aggregates.observed_without + excluded.observed_without,
3296 observed_with = usage_instance_aggregates.observed_with + excluded.observed_with,
3297 modeled_without = usage_instance_aggregates.modeled_without + excluded.modeled_without,
3298 modeled_with = usage_instance_aggregates.modeled_with + excluded.modeled_with,
3299 deduped_modeled_without = usage_instance_aggregates.deduped_modeled_without + excluded.deduped_modeled_without,
3300 deduped_modeled_with = usage_instance_aggregates.deduped_modeled_with + excluded.deduped_modeled_with,
3301 repeated_baselines = usage_instance_aggregates.repeated_baselines + excluded.repeated_baselines,
3302 observed_file_read_replacements = usage_instance_aggregates.observed_file_read_replacements + excluded.observed_file_read_replacements,
3303 modeled_file_reads_avoided = usage_instance_aggregates.modeled_file_reads_avoided + excluded.modeled_file_reads_avoided",
3304 aggregate_params!(instance_row_id, dimension_id, delta),
3305 )
3306 .map_err(aggregate_write_error)?;
3307 Ok(())
3308}
3309
3310fn upsert_global_daily(
3312 connection: &Connection,
3313 project: ProjectInstanceId,
3314 day: i64,
3315 dimension_id: i64,
3316 delta: AggregateCounters,
3317) -> DbResult<()> {
3318 connection.execute(
3319 "INSERT INTO usage_daily_aggregates VALUES(
3320 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)
3321 ON CONFLICT(project_instance_id, day_epoch, dimension_id) DO UPDATE SET
3322 calls = usage_daily_aggregates.calls + excluded.calls,
3323 estimated_without = usage_daily_aggregates.estimated_without + excluded.estimated_without,
3324 estimated_with = usage_daily_aggregates.estimated_with + excluded.estimated_with,
3325 observed_without = usage_daily_aggregates.observed_without + excluded.observed_without,
3326 observed_with = usage_daily_aggregates.observed_with + excluded.observed_with,
3327 modeled_without = usage_daily_aggregates.modeled_without + excluded.modeled_without,
3328 modeled_with = usage_daily_aggregates.modeled_with + excluded.modeled_with,
3329 deduped_modeled_without = usage_daily_aggregates.deduped_modeled_without + excluded.deduped_modeled_without,
3330 deduped_modeled_with = usage_daily_aggregates.deduped_modeled_with + excluded.deduped_modeled_with,
3331 repeated_baselines = usage_daily_aggregates.repeated_baselines + excluded.repeated_baselines,
3332 observed_file_read_replacements = usage_daily_aggregates.observed_file_read_replacements + excluded.observed_file_read_replacements,
3333 modeled_file_reads_avoided = usage_daily_aggregates.modeled_file_reads_avoided + excluded.modeled_file_reads_avoided",
3334 daily_aggregate_params!(project.as_bytes().as_slice(), day, dimension_id, delta),
3335 )
3336 .map_err(aggregate_write_error)?;
3337 Ok(())
3338}
3339
3340fn upsert_instance_daily(
3342 connection: &Connection,
3343 instance_row_id: i64,
3344 day: i64,
3345 dimension_id: i64,
3346 delta: AggregateCounters,
3347) -> DbResult<()> {
3348 connection.execute(
3349 "INSERT INTO usage_instance_daily_aggregates VALUES(
3350 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)
3351 ON CONFLICT(instance_row_id, day_epoch, dimension_id) DO UPDATE SET
3352 calls = usage_instance_daily_aggregates.calls + excluded.calls,
3353 estimated_without = usage_instance_daily_aggregates.estimated_without + excluded.estimated_without,
3354 estimated_with = usage_instance_daily_aggregates.estimated_with + excluded.estimated_with,
3355 observed_without = usage_instance_daily_aggregates.observed_without + excluded.observed_without,
3356 observed_with = usage_instance_daily_aggregates.observed_with + excluded.observed_with,
3357 modeled_without = usage_instance_daily_aggregates.modeled_without + excluded.modeled_without,
3358 modeled_with = usage_instance_daily_aggregates.modeled_with + excluded.modeled_with,
3359 deduped_modeled_without = usage_instance_daily_aggregates.deduped_modeled_without + excluded.deduped_modeled_without,
3360 deduped_modeled_with = usage_instance_daily_aggregates.deduped_modeled_with + excluded.deduped_modeled_with,
3361 repeated_baselines = usage_instance_daily_aggregates.repeated_baselines + excluded.repeated_baselines,
3362 observed_file_read_replacements = usage_instance_daily_aggregates.observed_file_read_replacements + excluded.observed_file_read_replacements,
3363 modeled_file_reads_avoided = usage_instance_daily_aggregates.modeled_file_reads_avoided + excluded.modeled_file_reads_avoided",
3364 daily_aggregate_params!(instance_row_id, day, dimension_id, delta),
3365 )
3366 .map_err(aggregate_write_error)?;
3367 Ok(())
3368}
3369
3370fn aggregate_write_error(error: rusqlite::Error) -> DbError {
3372 if matches!(
3373 &error,
3374 rusqlite::Error::SqliteFailure(sqlite_error, _)
3375 if sqlite_error.extended_code == rusqlite::ffi::SQLITE_CONSTRAINT_DATATYPE
3376 ) {
3377 DbError::TelemetryIntegerOverflow {
3378 field: AGGREGATE_COUNTER_FIELD,
3379 }
3380 } else {
3381 error.into()
3382 }
3383}
3384
3385#[derive(Debug, Eq, PartialEq)]
3387enum DailyEvictionCandidate {
3388 Global {
3390 project: Vec<u8>,
3392 day: i64,
3394 dimension_id: i64,
3396 },
3397 Instance {
3399 instance_row_id: i64,
3401 day: i64,
3403 dimension_id: i64,
3405 },
3406}
3407
3408impl DailyEvictionCandidate {
3409 fn ordering_key(&self) -> (i64, u8, Vec<u8>, i64) {
3411 match self {
3412 Self::Global {
3413 project,
3414 day,
3415 dimension_id,
3416 } => (*day, 0, project.clone(), *dimension_id),
3417 Self::Instance {
3418 instance_row_id,
3419 day,
3420 dimension_id,
3421 } => (
3422 *day,
3423 1,
3424 instance_row_id.to_be_bytes().to_vec(),
3425 *dimension_id,
3426 ),
3427 }
3428 }
3429}
3430
3431fn prepare_daily_capacity(
3433 connection: &Connection,
3434 project: ProjectInstanceId,
3435 instance_row_id: i64,
3436 day: i64,
3437 dimension_id: i64,
3438 policy: TelemetryRetentionPolicy,
3439) -> DbResult<()> {
3440 let global_exists = connection.query_row(
3441 "SELECT EXISTS(
3442 SELECT 1 FROM usage_daily_aggregates
3443 WHERE project_instance_id = ?1 AND day_epoch = ?2 AND dimension_id = ?3
3444 )",
3445 params![project.as_bytes().as_slice(), day, dimension_id],
3446 |row| row.get::<_, i64>(0),
3447 )?;
3448 let instance_exists = connection.query_row(
3449 "SELECT EXISTS(
3450 SELECT 1 FROM usage_instance_daily_aggregates
3451 WHERE instance_row_id = ?1 AND day_epoch = ?2 AND dimension_id = ?3
3452 )",
3453 params![instance_row_id, day, dimension_id],
3454 |row| row.get::<_, i64>(0),
3455 )?;
3456 let required = usize::from(global_exists == 0)
3457 .checked_add(usize::from(instance_exists == 0))
3458 .ok_or(DbError::TelemetryIntegerOverflow {
3459 field: "daily_rows",
3460 })?;
3461 if required == 0 {
3462 return Ok(());
3463 }
3464 let current = retention_counter(connection, RetentionCounter::DailyRows)?;
3465 if current > policy.max_daily_rows {
3466 return Err(DbError::TelemetryLimitInvalid {
3467 field: "daily_rows",
3468 value: current,
3469 });
3470 }
3471 let projected = current
3472 .checked_add(required)
3473 .ok_or(DbError::TelemetryIntegerOverflow {
3474 field: "daily_rows",
3475 })?;
3476 let evictions = projected.saturating_sub(policy.max_daily_rows);
3477 for _ in 0..evictions {
3478 let candidate = oldest_daily_eviction_candidate(
3479 connection,
3480 project,
3481 instance_row_id,
3482 day,
3483 dimension_id,
3484 )?
3485 .ok_or(DbError::TelemetryLimitInvalid {
3486 field: "max_daily_rows",
3487 value: policy.max_daily_rows,
3488 })?;
3489 delete_daily_eviction_candidate(connection, &candidate)?;
3490 }
3491 if evictions != 0 {
3492 decrement_retention_counter(connection, RetentionCounter::DailyRows, evictions)?;
3493 connection.execute(
3494 "UPDATE usage_retention_state SET label_history_complete = 0 WHERE singleton = 1",
3495 [],
3496 )?;
3497 }
3498 increment_retention_counter(connection, RetentionCounter::DailyRows, required)
3499}
3500
3501fn oldest_daily_eviction_candidate(
3503 connection: &Connection,
3504 project: ProjectInstanceId,
3505 instance_row_id: i64,
3506 day: i64,
3507 dimension_id: i64,
3508) -> DbResult<Option<DailyEvictionCandidate>> {
3509 let global = connection
3510 .query_row(
3511 "SELECT project_instance_id, day_epoch, dimension_id
3512 FROM usage_daily_aggregates
3513 WHERE NOT (
3514 project_instance_id = ?1 AND day_epoch = ?2 AND dimension_id = ?3
3515 )
3516 ORDER BY day_epoch, project_instance_id, dimension_id LIMIT 1",
3517 params![project.as_bytes().as_slice(), day, dimension_id],
3518 |row| {
3519 Ok(DailyEvictionCandidate::Global {
3520 project: row.get(0)?,
3521 day: row.get(1)?,
3522 dimension_id: row.get(2)?,
3523 })
3524 },
3525 )
3526 .optional()?;
3527 let instance = connection
3528 .query_row(
3529 "SELECT instance_row_id, day_epoch, dimension_id
3530 FROM usage_instance_daily_aggregates
3531 WHERE NOT (
3532 instance_row_id = ?1 AND day_epoch = ?2 AND dimension_id = ?3
3533 )
3534 ORDER BY day_epoch, instance_row_id, dimension_id LIMIT 1",
3535 params![instance_row_id, day, dimension_id],
3536 |row| {
3537 Ok(DailyEvictionCandidate::Instance {
3538 instance_row_id: row.get(0)?,
3539 day: row.get(1)?,
3540 dimension_id: row.get(2)?,
3541 })
3542 },
3543 )
3544 .optional()?;
3545 Ok(match (global, instance) {
3546 (Some(global), Some(instance)) => {
3547 if global.ordering_key() <= instance.ordering_key() {
3548 Some(global)
3549 } else {
3550 Some(instance)
3551 }
3552 }
3553 (Some(candidate), None) | (None, Some(candidate)) => Some(candidate),
3554 (None, None) => None,
3555 })
3556}
3557
3558fn delete_daily_eviction_candidate(
3560 connection: &Connection,
3561 candidate: &DailyEvictionCandidate,
3562) -> DbResult<()> {
3563 let deleted = match candidate {
3564 DailyEvictionCandidate::Global {
3565 project,
3566 day,
3567 dimension_id,
3568 } => connection.execute(
3569 "DELETE FROM usage_daily_aggregates
3570 WHERE project_instance_id = ?1 AND day_epoch = ?2 AND dimension_id = ?3",
3571 params![project, day, dimension_id],
3572 )?,
3573 DailyEvictionCandidate::Instance {
3574 instance_row_id,
3575 day,
3576 dimension_id,
3577 } => connection.execute(
3578 "DELETE FROM usage_instance_daily_aggregates
3579 WHERE instance_row_id = ?1 AND day_epoch = ?2 AND dimension_id = ?3",
3580 params![instance_row_id, day, dimension_id],
3581 )?,
3582 };
3583 if deleted != 1 {
3584 return Err(DbError::TelemetryIntegerOverflow {
3585 field: "daily_rows",
3586 });
3587 }
3588 Ok(())
3589}
3590
3591fn touch_instance(
3593 connection: &Connection,
3594 instance_row_id: i64,
3595 now: i64,
3596 policy: TelemetryRetentionPolicy,
3597) -> DbResult<()> {
3598 let previous = connection.query_row(
3599 "SELECT last_seen_at_epoch FROM usage_instances WHERE instance_row_id = ?1",
3600 [instance_row_id],
3601 |row| row.get::<_, i64>(0),
3602 )?;
3603 let tolerance = to_i64(
3604 "future_clock_tolerance_seconds",
3605 policy.future_clock_tolerance_seconds,
3606 )?;
3607 let anomaly = previous
3608 > now
3609 .checked_add(tolerance)
3610 .ok_or(DbError::TelemetryIntegerOverflow {
3611 field: "future_clock_tolerance",
3612 })?;
3613 let observed = previous.max(now);
3614 connection.execute(
3615 "UPDATE usage_instances
3616 SET last_seen_at_epoch = ?2,
3617 clock_anomaly = CASE WHEN ?3 = 1 THEN 1 ELSE clock_anomaly END
3618 WHERE instance_row_id = ?1",
3619 params![instance_row_id, observed, i64::from(anomaly)],
3620 )?;
3621 if anomaly {
3622 connection.execute(
3623 "UPDATE usage_retention_state SET clock_anomaly = 1 WHERE singleton = 1",
3624 [],
3625 )?;
3626 }
3627 Ok(())
3628}
3629
3630fn expire_idle_instances(
3632 connection: &Connection,
3633 project: ProjectInstanceId,
3634 current: UsageInstanceId,
3635 policy: TelemetryRetentionPolicy,
3636 now: i64,
3637) -> DbResult<usize> {
3638 let cutoff = epoch_cutoff(
3639 now,
3640 policy.max_active_idle_seconds,
3641 "max_active_idle_seconds",
3642 )?;
3643 let mut statement = connection.prepare_cached(
3644 "SELECT instance_row_id FROM usage_instances
3645 WHERE project_instance_id = ?1 AND state = ?2
3646 AND runtime_instance_id <> ?3 AND last_seen_at_epoch < ?4
3647 ORDER BY last_seen_at_epoch, instance_row_id LIMIT ?5",
3648 )?;
3649 let rows = statement.query_map(
3650 params![
3651 project.as_bytes().as_slice(),
3652 INSTANCE_ACTIVE,
3653 current.as_bytes().as_slice(),
3654 cutoff,
3655 to_i64("prune_batch_rows", policy.prune_batch_rows)?,
3656 ],
3657 |row| row.get::<_, i64>(0),
3658 )?;
3659 let ids = rows.collect::<Result<Vec<_>, _>>()?;
3660 drop(statement);
3661 for row_id in &ids {
3662 connection.execute(
3663 "UPDATE usage_instances
3664 SET state = ?2, sealed_at_epoch = last_seen_at_epoch
3665 WHERE instance_row_id = ?1 AND state = ?3",
3666 params![row_id, INSTANCE_EXPIRED, INSTANCE_ACTIVE],
3667 )?;
3668 delete_instance_baselines(connection, *row_id)?;
3669 }
3670 Ok(ids.len())
3671}
3672
3673fn prune_instances_once(
3675 connection: &Connection,
3676 policy: TelemetryRetentionPolicy,
3677 now: i64,
3678 reserve_instances: usize,
3679) -> DbResult<(usize, usize)> {
3680 let active_project = current_project(connection)?.as_bytes();
3681 let projected = retention_counter(connection, RetentionCounter::InstanceRows)?
3682 .checked_add(reserve_instances)
3683 .ok_or(DbError::TelemetryIntegerOverflow {
3684 field: "instance_rows",
3685 })?;
3686 let excess = projected.saturating_sub(policy.max_retained_instances);
3687 let cutoff = epoch_cutoff(
3688 now,
3689 policy.retained_instance_seconds,
3690 "retained_instance_seconds",
3691 )?;
3692 let limit = if excess == 0 {
3693 policy.prune_batch_rows
3694 } else {
3695 excess.min(policy.prune_batch_rows)
3696 };
3697 let mut statement = connection.prepare_cached(
3698 "SELECT instance_row_id, project_instance_id, runtime_instance_id, caller_label
3699 FROM usage_instances
3700 WHERE state IN (?1, ?2) AND (?3 > 0 OR last_seen_at_epoch < ?4)
3701 ORDER BY last_seen_at_epoch, instance_row_id LIMIT ?5",
3702 )?;
3703 let rows = statement.query_map(
3704 params![
3705 INSTANCE_SEALED,
3706 INSTANCE_EXPIRED,
3707 to_i64("excess_instance_rows", excess)?,
3708 cutoff,
3709 to_i64("prune_batch_rows", limit)?,
3710 ],
3711 |row| {
3712 Ok((
3713 row.get::<_, i64>(0)?,
3714 row.get::<_, Vec<u8>>(1)?,
3715 row.get::<_, Vec<u8>>(2)?,
3716 row.get::<_, Option<String>>(3)?,
3717 ))
3718 },
3719 )?;
3720 let candidates = rows.collect::<Result<Vec<_>, _>>()?;
3721 drop(statement);
3722 let mut raw_rows = 0usize;
3723 for (row_id, project, runtime, caller_label) in &candidates {
3724 let (raw, raw_bytes) = connection.query_row(
3725 "SELECT COUNT(*), COALESCE(SUM(logical_bytes), 0)
3726 FROM usage_events WHERE instance_row_id = ?1",
3727 [row_id],
3728 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
3729 )?;
3730 raw_rows = raw_rows
3731 .checked_add(count_usize("pruned_instance_raw_rows", raw)?)
3732 .ok_or(DbError::TelemetryIntegerOverflow {
3733 field: "pruned_instance_raw_rows",
3734 })?;
3735 let daily_rows = connection.query_row(
3736 "SELECT COUNT(*) FROM usage_instance_daily_aggregates WHERE instance_row_id = ?1",
3737 [row_id],
3738 |row| row.get::<_, i64>(0),
3739 )?;
3740 let tombstone_exists = connection.query_row(
3741 "SELECT EXISTS(
3742 SELECT 1 FROM usage_instance_tombstones
3743 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2
3744 )",
3745 params![project, runtime],
3746 |row| row.get::<_, i64>(0),
3747 )?;
3748 connection.execute(
3749 "INSERT INTO usage_instance_tombstones(
3750 project_instance_id, runtime_instance_id, retired_at_epoch
3751 ) VALUES(?1, ?2, ?3)
3752 ON CONFLICT(project_instance_id, runtime_instance_id) DO UPDATE SET
3753 retired_at_epoch = excluded.retired_at_epoch",
3754 params![project, runtime, now],
3755 )?;
3756 if tombstone_exists == 0 {
3757 increment_retention_counter(connection, RetentionCounter::InstanceTombstoneRows, 1)?;
3758 }
3759 if let Some(label) = caller_label {
3760 connection.execute(
3761 "UPDATE usage_labels SET detail_complete = 0
3762 WHERE project_instance_id = ?1 AND caller_label = ?2",
3763 params![project, label],
3764 )?;
3765 upsert_label_tombstone(connection, project, label, now, Some(runtime.as_slice()))?;
3766 }
3767 delete_instance_baselines(connection, *row_id)?;
3768 connection.execute(
3769 "DELETE FROM usage_instances WHERE instance_row_id = ?1",
3770 [row_id],
3771 )?;
3772 decrement_retention_counter(connection, RetentionCounter::InstanceRows, 1)?;
3773 decrement_retention_counter(
3774 connection,
3775 RetentionCounter::RawRows,
3776 count_usize("instance_raw_rows", raw)?,
3777 )?;
3778 decrement_retention_counter(
3779 connection,
3780 RetentionCounter::RawLogicalBytes,
3781 count_usize("instance_raw_logical_bytes", raw_bytes)?,
3782 )?;
3783 decrement_retention_counter(
3784 connection,
3785 RetentionCounter::DailyRows,
3786 count_usize("instance_daily_rows", daily_rows)?,
3787 )?;
3788 if project.as_slice() != active_project.as_slice() {
3789 let project_still_retained = connection.query_row(
3790 "SELECT EXISTS(
3791 SELECT 1 FROM usage_instances WHERE project_instance_id = ?1 LIMIT 1
3792 )",
3793 [project],
3794 |row| row.get::<_, i64>(0),
3795 )?;
3796 if project_still_retained == 0 {
3797 connection.execute(
3798 "DELETE FROM usage_global_aggregates WHERE project_instance_id = ?1",
3799 [project],
3800 )?;
3801 }
3802 }
3803 }
3804 Ok((candidates.len(), raw_rows))
3805}
3806
3807fn raw_pressure(
3809 connection: &Connection,
3810 policy: TelemetryRetentionPolicy,
3811 now: i64,
3812) -> DbResult<(usize, usize, usize)> {
3813 let cutoff = epoch_cutoff(now, policy.max_raw_age_seconds, "max_raw_age_seconds")?;
3814 let old_rows = connection.query_row(
3815 "SELECT EXISTS(
3816 SELECT 1 FROM usage_events
3817 WHERE created_at_epoch < ?1
3818 ORDER BY created_at_epoch, id LIMIT 1
3819 )",
3820 [cutoff],
3821 |row| row.get::<_, i64>(0),
3822 )?;
3823 Ok((
3824 retention_counter(connection, RetentionCounter::RawRows)?,
3825 retention_counter(connection, RetentionCounter::RawLogicalBytes)?,
3826 usize::from(old_rows != 0),
3827 ))
3828}
3829
3830fn prune_raw_once(
3832 connection: &Connection,
3833 policy: TelemetryRetentionPolicy,
3834 now: i64,
3835) -> DbResult<usize> {
3836 let (rows, bytes, old_rows) = raw_pressure(connection, policy, now)?;
3837 if rows <= policy.max_raw_rows && bytes <= policy.max_raw_logical_bytes && old_rows == 0 {
3838 return Ok(0);
3839 }
3840 let cutoff = epoch_cutoff(now, policy.max_raw_age_seconds, "max_raw_age_seconds")?;
3841 let mut statement = connection.prepare_cached(
3842 "SELECT created_at_epoch, logical_bytes FROM usage_events
3843 ORDER BY created_at_epoch, id LIMIT ?1",
3844 )?;
3845 let selected = statement.query_map(
3846 [to_i64("prune_batch_rows", policy.prune_batch_rows)?],
3847 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
3848 )?;
3849 let selected = selected.collect::<Result<Vec<_>, _>>()?;
3850 drop(statement);
3851 let mut remaining_rows = rows;
3852 let mut remaining_bytes = bytes;
3853 let mut delete_count = 0usize;
3854 let mut logical_bytes = 0usize;
3855 for (created_at, event_bytes) in selected {
3856 if created_at >= cutoff
3857 && remaining_rows <= policy.max_raw_rows
3858 && remaining_bytes <= policy.max_raw_logical_bytes
3859 {
3860 break;
3861 }
3862 let event_bytes = count_usize("raw_logical_bytes", event_bytes)?;
3863 remaining_rows = remaining_rows
3864 .checked_sub(1)
3865 .ok_or(DbError::TelemetryIntegerOverflow { field: "raw_rows" })?;
3866 remaining_bytes =
3867 remaining_bytes
3868 .checked_sub(event_bytes)
3869 .ok_or(DbError::TelemetryIntegerOverflow {
3870 field: "raw_logical_bytes",
3871 })?;
3872 delete_count = delete_count
3873 .checked_add(1)
3874 .ok_or(DbError::TelemetryIntegerOverflow { field: "raw_rows" })?;
3875 logical_bytes =
3876 logical_bytes
3877 .checked_add(event_bytes)
3878 .ok_or(DbError::TelemetryIntegerOverflow {
3879 field: "raw_logical_bytes",
3880 })?;
3881 }
3882 if delete_count != 0 {
3883 let limit = to_i64("pruned_raw_rows", delete_count)?;
3884 connection.execute(
3885 "UPDATE usage_labels SET detail_complete = 0
3886 WHERE (project_instance_id, caller_label) IN (
3887 SELECT DISTINCT i.project_instance_id, i.caller_label
3888 FROM usage_events AS e
3889 JOIN usage_instances AS i USING(instance_row_id)
3890 WHERE i.caller_label IS NOT NULL
3891 AND e.id IN (
3892 SELECT id FROM usage_events
3893 ORDER BY created_at_epoch, id LIMIT ?1
3894 )
3895 )",
3896 [limit],
3897 )?;
3898 connection.execute(
3899 "UPDATE usage_instances SET raw_detail_complete = 0
3900 WHERE instance_row_id IN (
3901 SELECT instance_row_id FROM usage_events
3902 ORDER BY created_at_epoch, id LIMIT ?1
3903 )",
3904 [limit],
3905 )?;
3906 let deleted = connection.execute(
3907 "DELETE FROM usage_events WHERE id IN (
3908 SELECT id FROM usage_events
3909 ORDER BY created_at_epoch, id LIMIT ?1
3910 )",
3911 [limit],
3912 )?;
3913 if deleted != delete_count {
3914 return Err(DbError::TelemetryIntegerOverflow { field: "raw_rows" });
3915 }
3916 decrement_retention_counter(connection, RetentionCounter::RawRows, delete_count)?;
3917 decrement_retention_counter(connection, RetentionCounter::RawLogicalBytes, logical_bytes)?;
3918 connection.execute(
3919 "UPDATE usage_retention_state SET raw_detail_complete = 0 WHERE singleton = 1",
3920 [],
3921 )?;
3922 }
3923 Ok(delete_count)
3924}
3925
3926fn prune_daily_once(
3928 connection: &Connection,
3929 policy: TelemetryRetentionPolicy,
3930 now: i64,
3931) -> DbResult<usize> {
3932 let days = policy
3933 .retained_trend_days
3934 .checked_mul(SECONDS_PER_DAY as u64)
3935 .ok_or(DbError::TelemetryIntegerOverflow {
3936 field: "retained_trend_seconds",
3937 })?;
3938 let cutoff = epoch_cutoff(now, days, "retained_trend_seconds")?;
3939 let limit = to_i64("prune_batch_rows", policy.prune_batch_rows)?;
3940 let global = connection.execute(
3941 "DELETE FROM usage_daily_aggregates
3942 WHERE (project_instance_id, day_epoch, dimension_id) IN (
3943 SELECT project_instance_id, day_epoch, dimension_id
3944 FROM usage_daily_aggregates WHERE day_epoch < ?1
3945 ORDER BY day_epoch, project_instance_id, dimension_id LIMIT ?2
3946 )",
3947 params![cutoff, limit],
3948 )?;
3949 let remaining = policy.prune_batch_rows.saturating_sub(global);
3950 let instance = if remaining == 0 {
3951 0
3952 } else {
3953 connection.execute(
3954 "DELETE FROM usage_instance_daily_aggregates
3955 WHERE (instance_row_id, day_epoch, dimension_id) IN (
3956 SELECT instance_row_id, day_epoch, dimension_id
3957 FROM usage_instance_daily_aggregates WHERE day_epoch < ?1
3958 ORDER BY day_epoch, instance_row_id, dimension_id LIMIT ?2
3959 )",
3960 params![cutoff, to_i64("daily_prune_remaining", remaining)?],
3961 )?
3962 };
3963 let deleted = global
3964 .checked_add(instance)
3965 .ok_or(DbError::TelemetryIntegerOverflow {
3966 field: "pruned_daily_rows",
3967 })?;
3968 if deleted != 0 {
3969 decrement_retention_counter(connection, RetentionCounter::DailyRows, deleted)?;
3970 }
3971 Ok(deleted)
3972}
3973
3974fn prune_labels_once(
3976 connection: &Connection,
3977 policy: TelemetryRetentionPolicy,
3978 now: i64,
3979) -> DbResult<usize> {
3980 let cutoff = epoch_cutoff(now, policy.retained_label_seconds, "retained_label_seconds")?;
3981 let label = connection
3982 .query_row(
3983 "SELECT project_instance_id, caller_label FROM usage_labels AS labels
3984 WHERE last_seen_at_epoch < ?1
3985 AND NOT EXISTS(
3986 SELECT 1 FROM usage_instances AS i
3987 WHERE i.project_instance_id = labels.project_instance_id
3988 AND i.caller_label = labels.caller_label
3989 AND i.state = ?2
3990 )
3991 AND (
3992 SELECT COUNT(*) FROM usage_instances AS i
3993 WHERE i.project_instance_id = labels.project_instance_id
3994 AND i.caller_label = labels.caller_label
3995 ) <= ?3
3996 ORDER BY last_seen_at_epoch, project_instance_id, caller_label LIMIT 1",
3997 params![
3998 cutoff,
3999 INSTANCE_ACTIVE,
4000 to_i64("prune_batch_rows", policy.prune_batch_rows)?,
4001 ],
4002 |row| Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, String>(1)?)),
4003 )
4004 .optional()?;
4005 let Some((project, label)) = label else {
4006 return Ok(0);
4007 };
4008 let runtime = connection
4009 .query_row(
4010 "SELECT runtime_instance_id FROM usage_instances
4011 WHERE project_instance_id = ?1 AND caller_label = ?2
4012 ORDER BY last_seen_at_epoch DESC, instance_row_id DESC LIMIT 1",
4013 params![project, label],
4014 |row| row.get::<_, Vec<u8>>(0),
4015 )
4016 .optional()?;
4017 upsert_label_tombstone(connection, &project, &label, now, runtime.as_deref())?;
4018 connection.execute(
4019 "UPDATE usage_instances SET caller_label = NULL
4020 WHERE project_instance_id = ?1 AND caller_label = ?2 AND state <> ?3",
4021 params![project, label, INSTANCE_ACTIVE],
4022 )?;
4023 let deleted = connection.execute(
4024 "DELETE FROM usage_labels WHERE project_instance_id = ?1 AND caller_label = ?2",
4025 params![project, label],
4026 )?;
4027 decrement_retention_counter(connection, RetentionCounter::LabelRows, deleted)?;
4028 connection.execute(
4029 "UPDATE usage_retention_state SET label_history_complete = 0 WHERE singleton = 1",
4030 [],
4031 )?;
4032 Ok(deleted)
4033}
4034
4035fn prune_tombstones_once(
4037 connection: &Connection,
4038 policy: TelemetryRetentionPolicy,
4039 now: i64,
4040) -> DbResult<usize> {
4041 let cutoff = epoch_cutoff(
4042 now,
4043 policy.retained_tombstone_seconds,
4044 "retained_tombstone_seconds",
4045 )?;
4046 let label_rows = retention_counter(connection, RetentionCounter::LabelTombstoneRows)?;
4047 let mut statement = connection.prepare_cached(
4048 "SELECT expired_at_epoch FROM usage_label_tombstones
4049 ORDER BY expired_at_epoch, project_instance_id, caller_label LIMIT ?1",
4050 )?;
4051 let rows = statement.query_map(
4052 [to_i64("prune_batch_rows", policy.prune_batch_rows)?],
4053 |row| row.get::<_, i64>(0),
4054 )?;
4055 let mut remaining_label_rows = label_rows;
4056 let mut label_limit = 0usize;
4057 for expired_at in rows {
4058 let expired_at = expired_at?;
4059 if expired_at >= cutoff && remaining_label_rows <= policy.max_label_tombstones {
4060 break;
4061 }
4062 remaining_label_rows =
4063 remaining_label_rows
4064 .checked_sub(1)
4065 .ok_or(DbError::TelemetryIntegerOverflow {
4066 field: "label_tombstone_rows",
4067 })?;
4068 label_limit = label_limit
4069 .checked_add(1)
4070 .ok_or(DbError::TelemetryIntegerOverflow {
4071 field: "label_tombstone_rows",
4072 })?;
4073 }
4074 drop(statement);
4075 let label = connection.execute(
4076 "DELETE FROM usage_label_tombstones
4077 WHERE (project_instance_id, caller_label) IN (
4078 SELECT project_instance_id, caller_label FROM usage_label_tombstones
4079 ORDER BY expired_at_epoch, project_instance_id, caller_label LIMIT ?1
4080 )",
4081 [to_i64("label_tombstone_prune_rows", label_limit)?],
4082 )?;
4083 if label != label_limit {
4084 return Err(DbError::TelemetryIntegerOverflow {
4085 field: "label_tombstone_rows",
4086 });
4087 }
4088 let remaining = policy.prune_batch_rows.saturating_sub(label);
4089 let instance = if remaining == 0 {
4090 0
4091 } else {
4092 let instance_rows = retention_counter(connection, RetentionCounter::InstanceTombstoneRows)?;
4093 let mut statement = connection.prepare_cached(
4094 "SELECT retired_at_epoch FROM usage_instance_tombstones
4095 ORDER BY retired_at_epoch, project_instance_id, runtime_instance_id LIMIT ?1",
4096 )?;
4097 let rows = statement
4098 .query_map([to_i64("tombstone_prune_remaining", remaining)?], |row| {
4099 row.get::<_, i64>(0)
4100 })?;
4101 let mut remaining_instance_rows = instance_rows;
4102 let mut instance_limit = 0usize;
4103 for retired_at in rows {
4104 let retired_at = retired_at?;
4105 if retired_at >= cutoff && remaining_instance_rows <= policy.max_instance_tombstones {
4106 break;
4107 }
4108 remaining_instance_rows = remaining_instance_rows.checked_sub(1).ok_or(
4109 DbError::TelemetryIntegerOverflow {
4110 field: "instance_tombstone_rows",
4111 },
4112 )?;
4113 instance_limit =
4114 instance_limit
4115 .checked_add(1)
4116 .ok_or(DbError::TelemetryIntegerOverflow {
4117 field: "instance_tombstone_rows",
4118 })?;
4119 }
4120 drop(statement);
4121 let deleted = connection.execute(
4122 "DELETE FROM usage_instance_tombstones
4123 WHERE (project_instance_id, runtime_instance_id) IN (
4124 SELECT project_instance_id, runtime_instance_id
4125 FROM usage_instance_tombstones
4126 ORDER BY retired_at_epoch, project_instance_id, runtime_instance_id LIMIT ?1
4127 )",
4128 [to_i64("instance_tombstone_prune_rows", instance_limit)?],
4129 )?;
4130 if deleted != instance_limit {
4131 return Err(DbError::TelemetryIntegerOverflow {
4132 field: "instance_tombstone_rows",
4133 });
4134 }
4135 deleted
4136 };
4137 if label != 0 {
4138 decrement_retention_counter(connection, RetentionCounter::LabelTombstoneRows, label)?;
4139 }
4140 if instance != 0 {
4141 decrement_retention_counter(
4142 connection,
4143 RetentionCounter::InstanceTombstoneRows,
4144 instance,
4145 )?;
4146 }
4147 label
4148 .checked_add(instance)
4149 .ok_or(DbError::TelemetryIntegerOverflow {
4150 field: "evicted_tombstones",
4151 })
4152}
4153
4154fn reconcile_retention_counters(connection: &Connection) -> DbResult<()> {
4157 let counts = connection.query_row(
4158 "SELECT
4159 (SELECT COUNT(*) FROM usage_events),
4160 (SELECT COALESCE(SUM(logical_bytes), 0) FROM usage_events),
4161 (SELECT COUNT(*) FROM usage_instance_baselines),
4162 (SELECT COALESCE(SUM(witness_logical_bytes), 0)
4163 FROM usage_instance_baselines),
4164 (SELECT COUNT(*) FROM usage_bucket_dimensions),
4165 (SELECT COUNT(*) FROM usage_instances),
4166 (SELECT COUNT(*) FROM usage_labels),
4167 (SELECT COUNT(*) FROM usage_daily_aggregates)
4168 + (SELECT COUNT(*) FROM usage_instance_daily_aggregates),
4169 (SELECT COUNT(*) FROM usage_label_tombstones),
4170 (SELECT COUNT(*) FROM usage_instance_tombstones)",
4171 [],
4172 |row| {
4173 Ok((
4174 row.get::<_, i64>(0)?,
4175 row.get::<_, i64>(1)?,
4176 row.get::<_, i64>(2)?,
4177 row.get::<_, i64>(3)?,
4178 row.get::<_, i64>(4)?,
4179 row.get::<_, i64>(5)?,
4180 row.get::<_, i64>(6)?,
4181 row.get::<_, i64>(7)?,
4182 row.get::<_, i64>(8)?,
4183 row.get::<_, i64>(9)?,
4184 ))
4185 },
4186 )?;
4187 for (field, value) in [
4188 ("raw_rows", counts.0),
4189 ("raw_logical_bytes", counts.1),
4190 ("baseline_rows", counts.2),
4191 ("baseline_logical_bytes", counts.3),
4192 ("dimension_rows", counts.4),
4193 ("instance_rows", counts.5),
4194 ("label_rows", counts.6),
4195 ("daily_rows", counts.7),
4196 ("label_tombstone_rows", counts.8),
4197 ("instance_tombstone_rows", counts.9),
4198 ] {
4199 let _ = count_usize(field, value)?;
4200 }
4201 connection.execute(
4202 "UPDATE usage_retention_state
4203 SET raw_rows = ?1,
4204 raw_logical_bytes = ?2,
4205 baseline_rows = ?3,
4206 baseline_logical_bytes = ?4,
4207 dimension_rows = ?5,
4208 instance_rows = ?6,
4209 label_rows = ?7,
4210 daily_rows = ?8,
4211 label_tombstone_rows = ?9,
4212 instance_tombstone_rows = ?10
4213 WHERE singleton = 1",
4214 params![
4215 counts.0, counts.1, counts.2, counts.3, counts.4, counts.5, counts.6, counts.7,
4216 counts.8, counts.9,
4217 ],
4218 )?;
4219 Ok(())
4220}
4221
4222fn aged_maintenance_pending(
4224 connection: &Connection,
4225 policy: TelemetryRetentionPolicy,
4226 now: i64,
4227) -> DbResult<bool> {
4228 let instance_cutoff = epoch_cutoff(
4229 now,
4230 policy.retained_instance_seconds,
4231 "retained_instance_seconds",
4232 )?;
4233 let trend_seconds = policy
4234 .retained_trend_days
4235 .checked_mul(SECONDS_PER_DAY as u64)
4236 .ok_or(DbError::TelemetryIntegerOverflow {
4237 field: "retained_trend_seconds",
4238 })?;
4239 let daily_cutoff = epoch_cutoff(now, trend_seconds, "retained_trend_seconds")?;
4240 let label_cutoff = epoch_cutoff(now, policy.retained_label_seconds, "retained_label_seconds")?;
4241 let tombstone_cutoff = epoch_cutoff(
4242 now,
4243 policy.retained_tombstone_seconds,
4244 "retained_tombstone_seconds",
4245 )?;
4246 let pending = connection.query_row(
4247 "SELECT
4248 EXISTS(
4249 SELECT 1 FROM usage_instances
4250 WHERE state IN (?1, ?2) AND last_seen_at_epoch < ?3 LIMIT 1
4251 ) OR EXISTS(
4252 SELECT 1 FROM usage_daily_aggregates
4253 WHERE day_epoch < ?4 LIMIT 1
4254 ) OR EXISTS(
4255 SELECT 1 FROM usage_instance_daily_aggregates
4256 WHERE day_epoch < ?4 LIMIT 1
4257 ) OR EXISTS(
4258 SELECT 1 FROM usage_labels AS labels
4259 WHERE last_seen_at_epoch < ?5
4260 AND NOT EXISTS(
4261 SELECT 1 FROM usage_instances AS i
4262 WHERE i.project_instance_id = labels.project_instance_id
4263 AND i.caller_label = labels.caller_label
4264 AND i.state = ?6
4265 )
4266 AND (
4267 SELECT COUNT(*) FROM usage_instances AS i
4268 WHERE i.project_instance_id = labels.project_instance_id
4269 AND i.caller_label = labels.caller_label
4270 ) <= ?7
4271 LIMIT 1
4272 ) OR EXISTS(
4273 SELECT 1 FROM usage_label_tombstones
4274 WHERE expired_at_epoch < ?8 LIMIT 1
4275 ) OR EXISTS(
4276 SELECT 1 FROM usage_instance_tombstones
4277 WHERE retired_at_epoch < ?8 LIMIT 1
4278 )",
4279 params![
4280 INSTANCE_SEALED,
4281 INSTANCE_EXPIRED,
4282 instance_cutoff,
4283 daily_cutoff,
4284 label_cutoff,
4285 INSTANCE_ACTIVE,
4286 to_i64("prune_batch_rows", policy.prune_batch_rows)?,
4287 tombstone_cutoff,
4288 ],
4289 |row| row.get::<_, i64>(0),
4290 )?;
4291 Ok(pending != 0)
4292}
4293
4294#[allow(clippy::too_many_arguments)]
4295fn refresh_retention_state(
4297 connection: &Connection,
4298 policy: TelemetryRetentionPolicy,
4299 now: i64,
4300 pruned_raw: usize,
4301 pruned_instances: usize,
4302 evicted_tombstones: usize,
4303 writes_added: usize,
4304) -> DbResult<()> {
4305 let counters = connection.query_row(
4306 "SELECT raw_rows, raw_logical_bytes, instance_rows, label_rows, daily_rows,
4307 label_tombstone_rows, instance_tombstone_rows,
4308 pruned_raw_rows, pruned_instance_rows, evicted_tombstones,
4309 writes_since_checkpoint
4310 FROM usage_retention_state WHERE singleton = 1",
4311 [],
4312 |row| {
4313 Ok((
4314 row.get::<_, i64>(0)?,
4315 row.get::<_, i64>(1)?,
4316 row.get::<_, i64>(2)?,
4317 row.get::<_, i64>(3)?,
4318 row.get::<_, i64>(4)?,
4319 row.get::<_, i64>(5)?,
4320 row.get::<_, i64>(6)?,
4321 row.get::<_, i64>(7)?,
4322 row.get::<_, i64>(8)?,
4323 row.get::<_, i64>(9)?,
4324 row.get::<_, i64>(10)?,
4325 ))
4326 },
4327 )?;
4328 let raw_rows = count_usize("raw_rows", counters.0)?;
4329 let raw_bytes = count_usize("raw_logical_bytes", counters.1)?;
4330 let instance_rows = count_usize("instance_rows", counters.2)?;
4331 let label_rows = count_usize("label_rows", counters.3)?;
4332 let daily_rows = count_usize("daily_rows", counters.4)?;
4333 let label_tombstones = count_usize("label_tombstone_rows", counters.5)?;
4334 let instance_tombstones = count_usize("instance_tombstone_rows", counters.6)?;
4335 let cutoff = epoch_cutoff(now, policy.max_raw_age_seconds, "max_raw_age_seconds")?;
4336 let old_raw = connection.query_row(
4337 "SELECT EXISTS(
4338 SELECT 1 FROM usage_events
4339 WHERE created_at_epoch < ?1 LIMIT 1
4340 )",
4341 [cutoff],
4342 |row| row.get::<_, i64>(0),
4343 )?;
4344 let oldest = connection
4345 .query_row(
4346 "SELECT created_at_epoch FROM usage_events
4347 ORDER BY created_at_epoch, id LIMIT 1",
4348 [],
4349 |row| row.get::<_, i64>(0),
4350 )
4351 .optional()?;
4352 let pruned_raw_total = checked_count_add("pruned_raw_rows", counters.7, pruned_raw)?;
4353 let pruned_instance_total =
4354 checked_count_add("pruned_instance_rows", counters.8, pruned_instances)?;
4355 let evicted_total = checked_count_add("evicted_tombstones", counters.9, evicted_tombstones)?;
4356 let writes = checked_count_add("writes_since_checkpoint", counters.10, writes_added)?;
4357 let maintenance_pending = raw_rows > policy.max_raw_rows
4358 || raw_bytes > policy.max_raw_logical_bytes
4359 || old_raw > 0
4360 || instance_rows > policy.max_retained_instances
4361 || label_rows > policy.max_retained_labels
4362 || daily_rows > policy.max_daily_rows
4363 || label_tombstones > policy.max_label_tombstones
4364 || instance_tombstones > policy.max_instance_tombstones
4365 || count_usize("writes_since_checkpoint", writes)? >= policy.checkpoint_write_interval
4366 || aged_maintenance_pending(connection, policy, now)?;
4367 connection.execute(
4368 "UPDATE usage_retention_state
4369 SET policy_version = ?1,
4370 logical_byte_version = ?2,
4371 pruned_raw_rows = ?3,
4372 pruned_instance_rows = ?4,
4373 evicted_tombstones = ?5,
4374 writes_since_checkpoint = ?6,
4375 last_maintenance_epoch = ?7,
4376 oldest_retained_epoch = ?8,
4377 maintenance_pending = ?9
4378 WHERE singleton = 1",
4379 params![
4380 i64::from(POLICY_VERSION),
4381 i64::from(LOGICAL_BYTE_VERSION),
4382 pruned_raw_total,
4383 pruned_instance_total,
4384 evicted_total,
4385 writes,
4386 now,
4387 oldest,
4388 i64::from(maintenance_pending),
4389 ],
4390 )?;
4391 Ok(())
4392}
4393
4394fn converge_retention(
4396 connection: &Connection,
4397 project: ProjectInstanceId,
4398 policy: TelemetryRetentionPolicy,
4399 now: i64,
4400) -> DbResult<()> {
4401 let mut pruned_raw = 0usize;
4402 loop {
4403 let (rows, bytes, old_rows) = raw_pressure(connection, policy, now)?;
4404 if rows <= policy.max_raw_rows && bytes <= policy.max_raw_logical_bytes && old_rows == 0 {
4405 break;
4406 }
4407 let deleted = prune_raw_once(connection, policy, now)?;
4408 if deleted == 0 {
4409 break;
4410 }
4411 pruned_raw = pruned_raw
4412 .checked_add(deleted)
4413 .ok_or(DbError::TelemetryIntegerOverflow {
4414 field: "pruned_raw_rows",
4415 })?;
4416 }
4417 let mut pruned_instances = 0usize;
4418 loop {
4419 if retention_counter(connection, RetentionCounter::InstanceRows)?
4420 <= policy.max_retained_instances
4421 {
4422 break;
4423 }
4424 let (deleted, raw) = prune_instances_once(connection, policy, now, 0)?;
4425 if deleted == 0 {
4426 break;
4427 }
4428 pruned_instances =
4429 pruned_instances
4430 .checked_add(deleted)
4431 .ok_or(DbError::TelemetryIntegerOverflow {
4432 field: "pruned_instance_rows",
4433 })?;
4434 pruned_raw = pruned_raw
4435 .checked_add(raw)
4436 .ok_or(DbError::TelemetryIntegerOverflow {
4437 field: "pruned_raw_rows",
4438 })?;
4439 }
4440 while prune_daily_once(connection, policy, now)? == policy.prune_batch_rows {}
4441 let mut evicted = 0usize;
4442 loop {
4443 let deleted = prune_tombstones_once(connection, policy, now)?;
4444 evicted = evicted
4445 .checked_add(deleted)
4446 .ok_or(DbError::TelemetryIntegerOverflow {
4447 field: "evicted_tombstones",
4448 })?;
4449 if deleted < policy.prune_batch_rows {
4450 break;
4451 }
4452 }
4453 refresh_retention_state(
4454 connection,
4455 policy,
4456 now,
4457 pruned_raw,
4458 pruned_instances,
4459 evicted,
4460 0,
4461 )?;
4462 retention_state_for_project(connection, project).map(|_| ())
4463}
4464
4465fn raw_event_select(extra_predicate: &str) -> String {
4467 format!(
4468 "SELECT COALESCE(i.caller_label, ''), e.command, e.path, e.query,
4469 e.estimated_tokens_without_projectatlas,
4470 e.estimated_tokens_with_projectatlas, e.estimated_tokens_saved,
4471 d.token_savings_bucket, d.provider, d.model, d.tokenizer_backend,
4472 d.accuracy, d.baseline_kind, d.confidence, e.calculation_trace,
4473 d.accounting_layer, d.estimate_method, d.denominator_kind,
4474 e.baseline_identity, e.baseline_fingerprint, d.dedupe_scope
4475 FROM usage_events AS e
4476 JOIN usage_instances AS i USING(instance_row_id)
4477 JOIN usage_bucket_dimensions AS d USING(dimension_id)
4478 WHERE i.project_instance_id = ?1 {extra_predicate}
4479 ORDER BY e.id"
4480 )
4481}
4482
4483fn map_usage_event(row: &rusqlite::Row<'_>) -> DbResult<UsageEvent> {
4485 Ok(UsageEvent {
4486 session_id: row.get(0)?,
4487 command: row.get(1)?,
4488 path: row.get(2)?,
4489 query: row.get(3)?,
4490 estimated_tokens_without_projectatlas: row
4491 .get::<_, Option<i64>>(4)?
4492 .map(|value| count_usize("estimated_tokens_without_projectatlas", value))
4493 .transpose()?,
4494 estimated_tokens_with_projectatlas: row
4495 .get::<_, Option<i64>>(5)?
4496 .map(|value| count_usize("estimated_tokens_with_projectatlas", value))
4497 .transpose()?,
4498 estimated_tokens_saved: row.get(6)?,
4499 token_savings_bucket: row.get(7)?,
4500 provider: row.get(8)?,
4501 model: row.get(9)?,
4502 tokenizer_backend: row.get(10)?,
4503 accuracy: row.get(11)?,
4504 baseline_kind: row.get(12)?,
4505 confidence: row.get(13)?,
4506 calculation_trace: row.get(14)?,
4507 accounting_layer: row.get(15)?,
4508 estimate_method: row.get(16)?,
4509 denominator_kind: row.get(17)?,
4510 baseline_identity: row.get(18)?,
4511 baseline_fingerprint: row.get(19)?,
4512 dedupe_scope: row.get(20)?,
4513 })
4514}
4515
4516fn load_overview_aggregates(
4518 connection: &Connection,
4519 project: ProjectInstanceId,
4520 caller_label: Option<&str>,
4521) -> DbResult<Vec<(DimensionValues, AggregateCounters)>> {
4522 let (sql, label) = if let Some(label) = caller_label {
4523 (
4524 "SELECT d.token_savings_bucket, d.provider, d.model, d.tokenizer_backend,
4525 d.accuracy, d.baseline_kind, d.confidence, d.accounting_layer,
4526 d.estimate_method, d.denominator_kind, d.dedupe_scope, d.overflow,
4527 a.calls, a.estimated_without, a.estimated_with,
4528 a.observed_without, a.observed_with, a.modeled_without, a.modeled_with,
4529 a.deduped_modeled_without, a.deduped_modeled_with,
4530 a.repeated_baselines, a.observed_file_read_replacements,
4531 a.modeled_file_reads_avoided
4532 FROM (
4533 SELECT aggregate.dimension_id,
4534 SUM(aggregate.calls) AS calls,
4535 SUM(aggregate.estimated_without) AS estimated_without,
4536 SUM(aggregate.estimated_with) AS estimated_with,
4537 SUM(aggregate.observed_without) AS observed_without,
4538 SUM(aggregate.observed_with) AS observed_with,
4539 SUM(aggregate.modeled_without) AS modeled_without,
4540 SUM(aggregate.modeled_with) AS modeled_with,
4541 SUM(aggregate.deduped_modeled_without) AS deduped_modeled_without,
4542 SUM(aggregate.deduped_modeled_with) AS deduped_modeled_with,
4543 SUM(aggregate.repeated_baselines) AS repeated_baselines,
4544 SUM(aggregate.observed_file_read_replacements)
4545 AS observed_file_read_replacements,
4546 SUM(aggregate.modeled_file_reads_avoided)
4547 AS modeled_file_reads_avoided
4548 FROM usage_instance_aggregates AS aggregate
4549 JOIN usage_instances AS instance USING(instance_row_id)
4550 WHERE instance.project_instance_id = ?1 AND instance.caller_label = ?2
4551 GROUP BY aggregate.dimension_id
4552 ) AS a
4553 JOIN usage_bucket_dimensions AS d USING(dimension_id)
4554 ORDER BY d.dimension_id",
4555 Some(label),
4556 )
4557 } else {
4558 (
4559 "SELECT d.token_savings_bucket, d.provider, d.model, d.tokenizer_backend,
4560 d.accuracy, d.baseline_kind, d.confidence, d.accounting_layer,
4561 d.estimate_method, d.denominator_kind, d.dedupe_scope, d.overflow,
4562 a.calls, a.estimated_without, a.estimated_with,
4563 a.observed_without, a.observed_with, a.modeled_without, a.modeled_with,
4564 a.deduped_modeled_without, a.deduped_modeled_with,
4565 a.repeated_baselines, a.observed_file_read_replacements,
4566 a.modeled_file_reads_avoided
4567 FROM usage_global_aggregates AS a
4568 JOIN usage_bucket_dimensions AS d USING(dimension_id)
4569 WHERE a.project_instance_id = ?1
4570 ORDER BY d.dimension_id",
4571 None,
4572 )
4573 };
4574 let mut statement = connection.prepare_cached(sql)?;
4575 let mut rows = if let Some(label) = label {
4576 statement.query(params![project.as_bytes().as_slice(), label])?
4577 } else {
4578 statement.query([project.as_bytes().as_slice()])?
4579 };
4580 let mut result = Vec::new();
4581 while let Some(row) = rows.next()? {
4582 result.push((read_dimension(row, 0)?, read_counters_offset(row, 12)?));
4583 }
4584 Ok(result)
4585}
4586
4587fn load_worktree_overview_aggregates(
4589 connection: &Connection,
4590 registration_id: Option<i64>,
4591 synchronized_only: bool,
4592) -> DbResult<Vec<(DimensionValues, AggregateCounters)>> {
4593 let predicate = if registration_id.is_some() {
4594 "a.registration_id = ?1"
4595 } else if synchronized_only {
4596 "a.source_kind = ?1"
4597 } else {
4598 return Err(DbError::WorktreeRegistrationRow {
4599 reason: "worktree aggregate scope is unbounded",
4600 });
4601 };
4602 let sql = format!(
4603 "SELECT d.token_savings_bucket, d.provider, d.model, d.tokenizer_backend,
4604 d.accuracy, d.baseline_kind, d.confidence, d.accounting_layer,
4605 d.estimate_method, d.denominator_kind, d.dedupe_scope, d.overflow,
4606 SUM(a.calls), SUM(a.estimated_without), SUM(a.estimated_with),
4607 SUM(a.observed_without), SUM(a.observed_with),
4608 SUM(a.modeled_without), SUM(a.modeled_with),
4609 SUM(a.deduped_modeled_without), SUM(a.deduped_modeled_with),
4610 SUM(a.repeated_baselines), SUM(a.observed_file_read_replacements),
4611 SUM(a.modeled_file_reads_avoided)
4612 FROM worktree_usage_aggregates AS a
4613 INDEXED BY idx_worktree_usage_aggregates_day_registration
4614 JOIN usage_bucket_dimensions AS d USING(dimension_id)
4615 WHERE a.day_epoch = -1 AND {predicate}
4616 GROUP BY a.dimension_id
4617 ORDER BY a.dimension_id"
4618 );
4619 let mut statement = connection.prepare_cached(&sql)?;
4620 let mut rows = if let Some(registration_id) = registration_id {
4621 statement.query([registration_id])?
4622 } else {
4623 statement.query([WORKTREE_USAGE_SYNCHRONIZED])?
4624 };
4625 let mut result = Vec::new();
4626 while let Some(row) = rows.next()? {
4627 result.push((read_dimension(row, 0)?, read_counters_offset(row, 12)?));
4628 }
4629 Ok(result)
4630}
4631
4632fn worktree_aggregates_exist(connection: &Connection) -> DbResult<bool> {
4634 connection
4635 .query_row(
4636 "SELECT EXISTS(
4637 SELECT 1
4638 FROM worktree_usage_aggregates
4639 INDEXED BY idx_worktree_usage_aggregates_day_registration
4640 WHERE day_epoch = -1
4641 LIMIT 1
4642 )",
4643 [],
4644 |row| row.get(0),
4645 )
4646 .map_err(DbError::from)
4647}
4648
4649fn aggregate_report_rows(
4651 rows: Vec<(DimensionValues, AggregateCounters)>,
4652) -> DbResult<(Vec<TokenBucketOverview>, TokenAccountingTotals, bool)> {
4653 let mut by_dimension = BTreeMap::<DimensionValues, AggregateCounters>::new();
4654 for (dimension, counters) in rows {
4655 let entry = by_dimension.entry(dimension).or_default();
4656 *entry = entry.checked_add(counters)?;
4657 }
4658 let mut totals = TokenAccountingTotals::default();
4659 let mut buckets = Vec::with_capacity(by_dimension.len());
4660 let mut average_directory_without = 0u128;
4661 let mut average_directory_with = 0u128;
4662 let mut average_policy_complete = true;
4663 for (dimension, counters) in by_dimension {
4664 if dimension.overflow
4665 && dimension.denominator_kind == OVERFLOW_DIMENSION
4666 && counters.modeled_without > 0
4667 {
4668 average_policy_complete = false;
4669 }
4670 totals.measured_tokens_saved = totals
4671 .measured_tokens_saved
4672 .checked_add(component_difference(
4673 counters.observed_without,
4674 counters.observed_with,
4675 ))
4676 .ok_or(DbError::TelemetryIntegerOverflow {
4677 field: "measured_tokens_saved",
4678 })?;
4679 totals.gross_modeled_tokens_avoided = totals
4680 .gross_modeled_tokens_avoided
4681 .checked_add(component_difference(
4682 counters.modeled_without,
4683 counters.modeled_with,
4684 ))
4685 .ok_or(DbError::TelemetryIntegerOverflow {
4686 field: "gross_modeled_tokens_avoided",
4687 })?;
4688 let deduped_modeled_delta = component_difference(
4689 counters.deduped_modeled_without,
4690 counters.deduped_modeled_with,
4691 );
4692 totals.deduped_modeled_tokens_avoided = totals
4693 .deduped_modeled_tokens_avoided
4694 .checked_add(deduped_modeled_delta)
4695 .ok_or(DbError::TelemetryIntegerOverflow {
4696 field: "deduped_modeled_tokens_avoided",
4697 })?;
4698 if dimension.denominator_kind == TOKEN_BASELINE_DIRECTORY_WALK {
4699 let emitted_with = count_u128("modeled_with", counters.modeled_with)?;
4700 let retained_without = if deduped_modeled_delta >= 0 {
4701 emitted_with.checked_add(deduped_modeled_delta.unsigned_abs())
4702 } else {
4703 emitted_with.checked_sub(deduped_modeled_delta.unsigned_abs())
4704 }
4705 .ok_or(DbError::TelemetryIntegerOverflow {
4706 field: "average_directory_without",
4707 })?;
4708 average_directory_without = average_directory_without
4709 .checked_add(retained_without)
4710 .ok_or(DbError::TelemetryIntegerOverflow {
4711 field: "average_directory_without",
4712 })?;
4713 average_directory_with = average_directory_with.checked_add(emitted_with).ok_or(
4714 DbError::TelemetryIntegerOverflow {
4715 field: "average_directory_with",
4716 },
4717 )?;
4718 } else {
4719 totals.average_modeled_tokens_avoided = totals
4720 .average_modeled_tokens_avoided
4721 .checked_add(deduped_modeled_delta)
4722 .ok_or(DbError::TelemetryIntegerOverflow {
4723 field: "average_modeled_tokens_avoided",
4724 })?;
4725 }
4726 totals.repeated_baselines_deduped = totals
4727 .repeated_baselines_deduped
4728 .checked_add(count_u128(
4729 "repeated_baselines",
4730 counters.repeated_baselines,
4731 )?)
4732 .ok_or(DbError::TelemetryIntegerOverflow {
4733 field: "repeated_baselines",
4734 })?;
4735 totals.observed_file_read_replacements = totals
4736 .observed_file_read_replacements
4737 .checked_add(count_u128(
4738 "observed_file_read_replacements",
4739 counters.observed_file_read_replacements,
4740 )?)
4741 .ok_or(DbError::TelemetryIntegerOverflow {
4742 field: "observed_file_read_replacements",
4743 })?;
4744 totals.modeled_file_reads_avoided = totals
4745 .modeled_file_reads_avoided
4746 .checked_add(count_u128(
4747 "modeled_file_reads_avoided",
4748 counters.modeled_file_reads_avoided,
4749 )?)
4750 .ok_or(DbError::TelemetryIntegerOverflow {
4751 field: "modeled_file_reads_avoided",
4752 })?;
4753 buckets.push(bucket_from_counters(dimension, counters)?);
4754 }
4755 let average_directory_without =
4756 average_modeled_baseline_tokens(TOKEN_BASELINE_DIRECTORY_WALK, average_directory_without);
4757 let average_directory_delta = i128::try_from(average_directory_without)
4758 .ok()
4759 .and_then(|without| {
4760 i128::try_from(average_directory_with)
4761 .ok()
4762 .and_then(|with| without.checked_sub(with))
4763 })
4764 .ok_or(DbError::TelemetryIntegerOverflow {
4765 field: "average_modeled_tokens_avoided",
4766 })?;
4767 totals.average_modeled_tokens_avoided = totals
4768 .average_modeled_tokens_avoided
4769 .checked_add(average_directory_delta)
4770 .ok_or(DbError::TelemetryIntegerOverflow {
4771 field: "average_modeled_tokens_avoided",
4772 })?;
4773 Ok((buckets, totals, average_policy_complete))
4774}
4775
4776fn bucket_from_counters(
4778 dimension: DimensionValues,
4779 counters: AggregateCounters,
4780) -> DbResult<TokenBucketOverview> {
4781 Ok(TokenBucketOverview::from_totals(
4782 dimension.token_savings_bucket,
4783 dimension.provider,
4784 dimension.model,
4785 dimension.tokenizer_backend,
4786 dimension.accuracy,
4787 dimension.baseline_kind,
4788 dimension.confidence,
4789 dimension.accounting_layer,
4790 dimension.estimate_method,
4791 dimension.denominator_kind,
4792 dimension.dedupe_scope,
4793 count_u128("calls", counters.calls)?,
4794 count_u128("estimated_without", counters.estimated_without)?,
4795 count_u128("estimated_with", counters.estimated_with)?,
4796 ))
4797}
4798
4799fn load_daily_aggregates(
4801 connection: &Connection,
4802 project: ProjectInstanceId,
4803 caller_label: Option<&str>,
4804 window: TokenTrendWindow,
4805) -> DbResult<Vec<(String, DimensionValues, AggregateCounters)>> {
4806 let period_expression = match window {
4807 TokenTrendWindow::Day => "strftime('%Y-%m-%d', a.day_epoch, 'unixepoch')",
4808 TokenTrendWindow::Week => "strftime('%Y-W%W', a.day_epoch, 'unixepoch')",
4809 TokenTrendWindow::Month => "strftime('%Y-%m', a.day_epoch, 'unixepoch')",
4810 TokenTrendWindow::Year => "strftime('%Y', a.day_epoch, 'unixepoch')",
4811 };
4812 let (table, join, predicate) = if caller_label.is_some() {
4813 (
4814 "usage_instance_daily_aggregates",
4815 "JOIN usage_instances AS i USING(instance_row_id)",
4816 "i.project_instance_id = ?1 AND i.caller_label = ?2",
4817 )
4818 } else {
4819 ("usage_daily_aggregates", "", "a.project_instance_id = ?1")
4820 };
4821 let sql = format!(
4822 "SELECT grouped.period,
4823 d.token_savings_bucket, d.provider, d.model, d.tokenizer_backend,
4824 d.accuracy, d.baseline_kind, d.confidence, d.accounting_layer,
4825 d.estimate_method, d.denominator_kind, d.dedupe_scope, d.overflow,
4826 grouped.calls, grouped.estimated_without, grouped.estimated_with,
4827 grouped.observed_without, grouped.observed_with,
4828 grouped.modeled_without, grouped.modeled_with,
4829 grouped.deduped_modeled_without, grouped.deduped_modeled_with,
4830 grouped.repeated_baselines, grouped.observed_file_read_replacements,
4831 grouped.modeled_file_reads_avoided
4832 FROM (
4833 SELECT {period_expression} AS period, a.dimension_id,
4834 SUM(a.calls) AS calls,
4835 SUM(a.estimated_without) AS estimated_without,
4836 SUM(a.estimated_with) AS estimated_with,
4837 SUM(a.observed_without) AS observed_without,
4838 SUM(a.observed_with) AS observed_with,
4839 SUM(a.modeled_without) AS modeled_without,
4840 SUM(a.modeled_with) AS modeled_with,
4841 SUM(a.deduped_modeled_without) AS deduped_modeled_without,
4842 SUM(a.deduped_modeled_with) AS deduped_modeled_with,
4843 SUM(a.repeated_baselines) AS repeated_baselines,
4844 SUM(a.observed_file_read_replacements)
4845 AS observed_file_read_replacements,
4846 SUM(a.modeled_file_reads_avoided) AS modeled_file_reads_avoided
4847 FROM {table} AS a {join}
4848 WHERE {predicate}
4849 GROUP BY period, a.dimension_id
4850 ) AS grouped
4851 JOIN usage_bucket_dimensions AS d USING(dimension_id)
4852 ORDER BY grouped.period, grouped.dimension_id"
4853 );
4854 let mut statement = connection.prepare(&sql)?;
4855 let mut rows = if let Some(label) = caller_label {
4856 statement.query(params![project.as_bytes().as_slice(), label])?
4857 } else {
4858 statement.query([project.as_bytes().as_slice()])?
4859 };
4860 let mut result = Vec::new();
4861 while let Some(row) = rows.next()? {
4862 result.push((
4863 row.get::<_, String>(0)?,
4864 read_dimension(row, 1)?,
4865 read_counters_offset(row, 13)?,
4866 ));
4867 }
4868 Ok(result)
4869}
4870
4871fn load_worktree_daily_aggregates(
4873 connection: &Connection,
4874 registration_id: Option<i64>,
4875 synchronized_only: bool,
4876 window: TokenTrendWindow,
4877) -> DbResult<Vec<(String, DimensionValues, AggregateCounters)>> {
4878 let period_expression = match window {
4879 TokenTrendWindow::Day => "strftime('%Y-%m-%d', a.day_epoch, 'unixepoch')",
4880 TokenTrendWindow::Week => "strftime('%Y-W%W', a.day_epoch, 'unixepoch')",
4881 TokenTrendWindow::Month => "strftime('%Y-%m', a.day_epoch, 'unixepoch')",
4882 TokenTrendWindow::Year => "strftime('%Y', a.day_epoch, 'unixepoch')",
4883 };
4884 let predicate = if registration_id.is_some() {
4885 "a.registration_id = ?1"
4886 } else if synchronized_only {
4887 "a.source_kind = ?1"
4888 } else {
4889 return Err(DbError::WorktreeRegistrationRow {
4890 reason: "worktree trend scope is unbounded",
4891 });
4892 };
4893 let sql = format!(
4894 "SELECT grouped.period,
4895 d.token_savings_bucket, d.provider, d.model, d.tokenizer_backend,
4896 d.accuracy, d.baseline_kind, d.confidence, d.accounting_layer,
4897 d.estimate_method, d.denominator_kind, d.dedupe_scope, d.overflow,
4898 grouped.calls, grouped.estimated_without, grouped.estimated_with,
4899 grouped.observed_without, grouped.observed_with,
4900 grouped.modeled_without, grouped.modeled_with,
4901 grouped.deduped_modeled_without, grouped.deduped_modeled_with,
4902 grouped.repeated_baselines, grouped.observed_file_read_replacements,
4903 grouped.modeled_file_reads_avoided
4904 FROM (
4905 SELECT {period_expression} AS period, a.dimension_id,
4906 SUM(a.calls) AS calls,
4907 SUM(a.estimated_without) AS estimated_without,
4908 SUM(a.estimated_with) AS estimated_with,
4909 SUM(a.observed_without) AS observed_without,
4910 SUM(a.observed_with) AS observed_with,
4911 SUM(a.modeled_without) AS modeled_without,
4912 SUM(a.modeled_with) AS modeled_with,
4913 SUM(a.deduped_modeled_without) AS deduped_modeled_without,
4914 SUM(a.deduped_modeled_with) AS deduped_modeled_with,
4915 SUM(a.repeated_baselines) AS repeated_baselines,
4916 SUM(a.observed_file_read_replacements)
4917 AS observed_file_read_replacements,
4918 SUM(a.modeled_file_reads_avoided) AS modeled_file_reads_avoided
4919 FROM worktree_usage_aggregates AS a
4920 INDEXED BY idx_worktree_usage_aggregates_day_registration
4921 WHERE a.day_epoch >= 0 AND {predicate}
4922 GROUP BY period, a.dimension_id
4923 ) AS grouped
4924 JOIN usage_bucket_dimensions AS d USING(dimension_id)
4925 ORDER BY grouped.period, grouped.dimension_id"
4926 );
4927 let mut statement = connection.prepare(&sql)?;
4928 let mut rows = if let Some(registration_id) = registration_id {
4929 statement.query([registration_id])?
4930 } else {
4931 statement.query([WORKTREE_USAGE_SYNCHRONIZED])?
4932 };
4933 let mut result = Vec::new();
4934 while let Some(row) = rows.next()? {
4935 result.push((
4936 row.get::<_, String>(0)?,
4937 read_dimension(row, 1)?,
4938 read_counters_offset(row, 13)?,
4939 ));
4940 }
4941 Ok(result)
4942}
4943
4944fn detail_availability(
4946 connection: &Connection,
4947 project: ProjectInstanceId,
4948 caller_label: Option<&str>,
4949) -> DbResult<UsageDetailAvailability> {
4950 let state = connection.query_row(
4951 "SELECT raw_detail_complete, dimension_detail_complete, label_history_complete
4952 FROM usage_retention_state WHERE singleton = 1",
4953 [],
4954 |row| {
4955 Ok((
4956 row.get::<_, i64>(0)?,
4957 row.get::<_, i64>(1)?,
4958 row.get::<_, i64>(2)?,
4959 ))
4960 },
4961 )?;
4962 let global_complete = bool_from_sql("raw_detail_complete", state.0)?
4963 && bool_from_sql("dimension_detail_complete", state.1)?
4964 && bool_from_sql("label_history_complete", state.2)?;
4965 let Some(label) = caller_label else {
4966 return Ok(if global_complete {
4967 UsageDetailAvailability::Retained
4968 } else {
4969 UsageDetailAvailability::Partial
4970 });
4971 };
4972 let instances = connection.query_row(
4973 "SELECT COUNT(*), MIN(raw_detail_complete)
4974 FROM usage_instances
4975 WHERE project_instance_id = ?1 AND caller_label = ?2",
4976 params![project.as_bytes().as_slice(), label],
4977 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, Option<i64>>(1)?)),
4978 )?;
4979 let label_detail = connection
4980 .query_row(
4981 "SELECT detail_complete FROM usage_labels
4982 WHERE project_instance_id = ?1 AND caller_label = ?2",
4983 params![project.as_bytes().as_slice(), label],
4984 |row| row.get::<_, i64>(0),
4985 )
4986 .optional()?;
4987 if instances.0 > 0 {
4988 return Ok(
4989 if global_complete && instances.1 == Some(1) && label_detail == Some(1) {
4990 UsageDetailAvailability::Retained
4991 } else {
4992 UsageDetailAvailability::Partial
4993 },
4994 );
4995 }
4996 let tombstone = connection.query_row(
4997 "SELECT EXISTS(
4998 SELECT 1 FROM usage_label_tombstones
4999 WHERE project_instance_id = ?1 AND caller_label = ?2
5000 )",
5001 params![project.as_bytes().as_slice(), label],
5002 |row| row.get::<_, i64>(0),
5003 )?;
5004 Ok(if tombstone != 0 || label_detail == Some(0) {
5005 UsageDetailAvailability::Expired
5006 } else {
5007 UsageDetailAvailability::Unavailable
5008 })
5009}
5010
5011fn read_dimension(row: &rusqlite::Row<'_>, offset: usize) -> DbResult<DimensionValues> {
5013 Ok(DimensionValues {
5014 token_savings_bucket: row.get(offset)?,
5015 provider: row.get(offset + 1)?,
5016 model: row.get(offset + 2)?,
5017 tokenizer_backend: row.get(offset + 3)?,
5018 accuracy: row.get(offset + 4)?,
5019 baseline_kind: row.get(offset + 5)?,
5020 confidence: row.get(offset + 6)?,
5021 accounting_layer: row.get(offset + 7)?,
5022 estimate_method: row.get(offset + 8)?,
5023 denominator_kind: row.get(offset + 9)?,
5024 dedupe_scope: row.get(offset + 10)?,
5025 overflow: bool_from_sql("usage_bucket_dimensions.overflow", row.get(offset + 11)?)?,
5026 })
5027}
5028
5029fn read_counters_offset(row: &rusqlite::Row<'_>, offset: usize) -> DbResult<AggregateCounters> {
5031 Ok(AggregateCounters {
5032 calls: row.get(offset)?,
5033 estimated_without: row.get(offset + 1)?,
5034 estimated_with: row.get(offset + 2)?,
5035 observed_without: row.get(offset + 3)?,
5036 observed_with: row.get(offset + 4)?,
5037 modeled_without: row.get(offset + 5)?,
5038 modeled_with: row.get(offset + 6)?,
5039 deduped_modeled_without: row.get(offset + 7)?,
5040 deduped_modeled_with: row.get(offset + 8)?,
5041 repeated_baselines: row.get(offset + 9)?,
5042 observed_file_read_replacements: row.get(offset + 10)?,
5043 modeled_file_reads_avoided: row.get(offset + 11)?,
5044 })
5045}
5046
5047fn logical_event_bytes(event: &UsageEvent, label: Option<&str>) -> DbResult<usize> {
5049 let identity = event.effective_baseline_identity();
5050 let fingerprint = event.effective_baseline_fingerprint();
5051 [
5052 label.unwrap_or_default(),
5053 event.command.as_str(),
5054 event.path.as_deref().unwrap_or_default(),
5055 event.query.as_deref().unwrap_or_default(),
5056 event.token_savings_bucket.as_str(),
5057 event.provider.as_str(),
5058 event.model.as_str(),
5059 event.tokenizer_backend.as_str(),
5060 event.accuracy.as_str(),
5061 event.baseline_kind.as_str(),
5062 event.confidence.as_str(),
5063 event.calculation_trace.as_str(),
5064 event.report_accounting_layer(),
5065 event.estimate_method.as_str(),
5066 event.report_denominator_kind(),
5067 identity.as_ref(),
5068 fingerprint.as_ref(),
5069 event.report_dedupe_scope(),
5070 ]
5071 .into_iter()
5072 .try_fold(96_usize, |total, value| {
5073 total
5074 .checked_add(value.len())
5075 .ok_or(DbError::TelemetryIntegerOverflow {
5076 field: "raw_logical_bytes",
5077 })
5078 })
5079}
5080
5081fn signed_components(value: i64) -> DbResult<(i64, i64)> {
5083 if value >= 0 {
5084 Ok((value, 0))
5085 } else {
5086 Ok((
5087 0,
5088 value
5089 .checked_neg()
5090 .ok_or(DbError::TelemetryIntegerOverflow {
5091 field: "signed_component",
5092 })?,
5093 ))
5094 }
5095}
5096
5097fn component_difference(without: i64, with: i64) -> i128 {
5099 i128::from(without) - i128::from(with)
5100}
5101
5102fn checked_count_add(field: &'static str, previous: i64, value: usize) -> DbResult<i64> {
5104 previous
5105 .checked_add(to_i64(field, value)?)
5106 .ok_or(DbError::TelemetryIntegerOverflow { field })
5107}
5108
5109fn duration_millis(field: &'static str, duration: Duration) -> DbResult<u64> {
5111 duration
5112 .as_secs()
5113 .checked_mul(1_000)
5114 .and_then(|milliseconds| milliseconds.checked_add(u64::from(duration.subsec_millis())))
5115 .ok_or(DbError::TelemetryIntegerOverflow { field })
5116}
5117
5118fn retention_counter(connection: &Connection, counter: RetentionCounter) -> DbResult<usize> {
5120 let value = connection.query_row(counter.select_sql(), [], |row| row.get::<_, i64>(0))?;
5121 count_usize(counter.field(), value)
5122}
5123
5124fn increment_retention_counter(
5126 connection: &Connection,
5127 counter: RetentionCounter,
5128 amount: usize,
5129) -> DbResult<()> {
5130 let value = retention_counter(connection, counter)?
5131 .checked_add(amount)
5132 .ok_or(DbError::TelemetryIntegerOverflow {
5133 field: counter.field(),
5134 })?;
5135 connection.execute(counter.update_sql(), [to_i64(counter.field(), value)?])?;
5136 Ok(())
5137}
5138
5139fn decrement_retention_counter(
5141 connection: &Connection,
5142 counter: RetentionCounter,
5143 amount: usize,
5144) -> DbResult<()> {
5145 let current = retention_counter(connection, counter)?;
5146 let value = current
5147 .checked_sub(amount)
5148 .ok_or(DbError::TelemetryIntegerOverflow {
5149 field: counter.field(),
5150 })?;
5151 connection.execute(counter.update_sql(), [to_i64(counter.field(), value)?])?;
5152 Ok(())
5153}
5154
5155fn bool_from_sql(field: &'static str, value: i64) -> DbResult<bool> {
5157 match value {
5158 0 => Ok(false),
5159 1 => Ok(true),
5160 _ => Err(DbError::InvalidEnum {
5161 field,
5162 value: value.to_string(),
5163 }),
5164 }
5165}
5166
5167fn option_usize_to_i64(field: &'static str, value: Option<usize>) -> DbResult<Option<i64>> {
5169 value.map(|value| to_i64(field, value)).transpose()
5170}
5171
5172fn option_isize_to_i64(field: &'static str, value: Option<isize>) -> DbResult<Option<i64>> {
5174 value
5175 .map(|value| {
5176 i64::try_from(value).map_err(|_source| DbError::TelemetryIntegerOverflow { field })
5177 })
5178 .transpose()
5179}
5180
5181fn to_i64(field: &'static str, value: impl TryInto<i64>) -> DbResult<i64> {
5183 value
5184 .try_into()
5185 .map_err(|_source| DbError::TelemetryIntegerOverflow { field })
5186}
5187
5188fn count_usize(field: &'static str, value: i64) -> DbResult<usize> {
5190 usize::try_from(value).map_err(|_source| DbError::TelemetryIntegerOverflow { field })
5191}
5192
5193fn count_u32(field: &'static str, value: i64) -> DbResult<u32> {
5195 u32::try_from(value).map_err(|_source| DbError::TelemetryIntegerOverflow { field })
5196}
5197
5198fn count_u64(field: &'static str, value: i64) -> DbResult<u64> {
5200 u64::try_from(value).map_err(|_source| DbError::TelemetryIntegerOverflow { field })
5201}
5202
5203fn count_u128(field: &'static str, value: i64) -> DbResult<u128> {
5205 u128::try_from(value).map_err(|_source| DbError::TelemetryIntegerOverflow { field })
5206}
5207
5208fn epoch_cutoff(now: i64, seconds: u64, field: &'static str) -> DbResult<i64> {
5210 let seconds = to_i64(field, seconds)?;
5211 Ok(now.checked_sub(seconds).unwrap_or(0))
5212}
5213
5214fn now_epoch_seconds() -> DbResult<i64> {
5216 let duration = SystemTime::now()
5217 .duration_since(UNIX_EPOCH)
5218 .map_err(|_source| DbError::TelemetryIntegerOverflow {
5219 field: "system_time",
5220 })?;
5221 to_i64("system_time", duration.as_secs())
5222}
5223
5224fn pragma_count(connection: &Connection, pragma: &str) -> DbResult<i64> {
5226 let sql = match pragma {
5227 "freelist_count" => "PRAGMA freelist_count",
5228 "page_count" => "PRAGMA page_count",
5229 "page_size" => "PRAGMA page_size",
5230 _ => {
5231 return Err(DbError::InvalidEnum {
5232 field: "telemetry_pragma",
5233 value: pragma.to_string(),
5234 });
5235 }
5236 };
5237 connection
5238 .query_row(sql, [], |row| row.get(0))
5239 .map_err(Into::into)
5240}
5241
5242fn synchronous_mode(value: i64) -> DbResult<&'static str> {
5244 match value {
5245 0 => Ok("off"),
5246 1 => Ok("normal"),
5247 2 => Ok("full"),
5248 3 => Ok("extra"),
5249 _ => Err(DbError::InvalidEnum {
5250 field: "pragma.synchronous",
5251 value: value.to_string(),
5252 }),
5253 }
5254}
5255
5256#[cfg(test)]
5257mod tests {
5258 use super::*;
5259 use crate::{AtlasStore, WorktreeAlias, WorktreeRegistrationState};
5260 use projectatlas_core::telemetry::{
5261 TOKEN_BASELINE_DIRECTORY_WALK, TOKEN_DEDUPE_SCOPE_EVENT, usage_from_estimates,
5262 usage_from_text,
5263 };
5264 use projectatlas_core::{Node, NodeKind, normalized_parent};
5265 use rusqlite::{Transaction, TransactionBehavior};
5266 use std::cell::RefCell;
5267 use std::error::Error;
5268 use std::fs;
5269 use std::io;
5270 use std::sync::mpsc::{Receiver, SyncSender, sync_channel};
5271
5272 thread_local! {
5273 static WORKTREE_TRACE: RefCell<Vec<String>> = const { RefCell::new(Vec::new()) };
5275 static WORKTREE_SNAPSHOT_EXPORT_BLOCKER: RefCell<Option<SnapshotExportBlocker>> = const { RefCell::new(None) };
5277 }
5278
5279 struct SnapshotExportBlocker {
5280 entered: SyncSender<()>,
5281 resume: Receiver<()>,
5282 }
5283
5284 fn record_worktree_statement(sql: &str) {
5285 WORKTREE_TRACE.with(|statements| statements.borrow_mut().push(sql.to_string()));
5286 }
5287
5288 #[allow(clippy::needless_pass_by_value)]
5290 fn record_worktree_event(event: rusqlite::trace::TraceEvent<'_>) {
5291 if let rusqlite::trace::TraceEvent::Stmt(_, sql) = event {
5292 record_worktree_statement(sql);
5293 }
5294 }
5295
5296 fn block_worktree_snapshot_aggregate_query(sql: &str) {
5297 if !sql.contains("SELECT -1 AS day_epoch") {
5298 return;
5299 }
5300 WORKTREE_SNAPSHOT_EXPORT_BLOCKER.with(|slot| {
5301 let Some(blocker) = slot.borrow_mut().take() else {
5302 return;
5303 };
5304 if blocker.entered.send(()).is_ok() {
5305 let _resume = blocker.resume.recv_timeout(Duration::from_secs(10));
5306 }
5307 });
5308 }
5309
5310 #[allow(clippy::needless_pass_by_value)]
5312 fn block_worktree_snapshot_aggregate_event(event: rusqlite::trace::TraceEvent<'_>) {
5313 if let rusqlite::trace::TraceEvent::Stmt(_, sql) = event {
5314 block_worktree_snapshot_aggregate_query(sql);
5315 }
5316 }
5317
5318 struct TestDatabase {
5319 temp: tempfile::TempDir,
5320 connection: Connection,
5321 project: ProjectInstanceId,
5322 }
5323
5324 struct ProductionTestDatabase {
5325 _temp: tempfile::TempDir,
5326 root: std::path::PathBuf,
5327 database_path: std::path::PathBuf,
5328 store: AtlasStore,
5329 project: ProjectInstanceId,
5330 }
5331
5332 fn test_database() -> Result<TestDatabase, Box<dyn Error>> {
5333 let temp = tempfile::tempdir()?;
5334 let database_path = temp.path().join("projectatlas.db");
5335 let connection = Connection::open(&database_path)?;
5336 crate::schema::initialize(&connection, None)?;
5337 let project = crate::project_identity::ensure_project_identity(&connection)?.0;
5338 Ok(TestDatabase {
5339 temp,
5340 connection,
5341 project,
5342 })
5343 }
5344
5345 fn production_database() -> Result<ProductionTestDatabase, Box<dyn Error>> {
5346 let temp = tempfile::tempdir()?;
5347 let root = temp.path().join("repository");
5348 let atlas = root.join(".projectatlas");
5349 fs::create_dir_all(&atlas)?;
5350 let database_path = atlas.join("projectatlas.db");
5351 let store = AtlasStore::open_for_project(&database_path, &root)?;
5352 let project = store
5353 .validated_project_instance_id
5354 .ok_or(DbError::ProjectInstanceIdentityMissing)?;
5355 Ok(ProductionTestDatabase {
5356 _temp: temp,
5357 root,
5358 database_path,
5359 store,
5360 project,
5361 })
5362 }
5363
5364 fn store_at(root: &std::path::Path) -> Result<AtlasStore, Box<dyn Error>> {
5365 let atlas = root.join(".projectatlas");
5366 fs::create_dir_all(&atlas)?;
5367 Ok(AtlasStore::open_for_project(
5368 &atlas.join("projectatlas.db"),
5369 root,
5370 )?)
5371 }
5372
5373 fn require(condition: bool, message: &str) -> Result<(), Box<dyn Error>> {
5375 if condition {
5376 Ok(())
5377 } else {
5378 Err(std::io::Error::other(message).into())
5379 }
5380 }
5381
5382 fn require_eq<T>(actual: &T, expected: &T, label: &str) -> Result<(), Box<dyn Error>>
5384 where
5385 T: std::fmt::Debug + PartialEq,
5386 {
5387 if actual == expected {
5388 Ok(())
5389 } else {
5390 Err(std::io::Error::other(format!(
5391 "{label} mismatch: expected {expected:?}, found {actual:?}"
5392 ))
5393 .into())
5394 }
5395 }
5396
5397 fn instance(byte: u8) -> Result<UsageInstanceId, Box<dyn Error>> {
5398 Ok(UsageInstanceId::from_bytes([byte; 16])?)
5399 }
5400
5401 fn event(label: &str, without: usize, with: usize) -> UsageEvent {
5402 let mut event = usage_from_estimates(
5403 label,
5404 "summary",
5405 Some("src/lib.rs".to_string()),
5406 None,
5407 without,
5408 with,
5409 );
5410 event.baseline_identity = "source:src/lib.rs".to_string();
5411 event.baseline_fingerprint = "source:src/lib.rs:v1".to_string();
5412 event
5413 }
5414
5415 fn directory_event(label: &str, without: usize, with: usize) -> UsageEvent {
5416 let mut event = event(label, without, with);
5417 event.command = "folders".to_string();
5418 event.denominator_kind = TOKEN_BASELINE_DIRECTORY_WALK.to_string();
5419 event.baseline_identity = "directory:src".to_string();
5420 event.baseline_fingerprint = "directory:src:v1".to_string();
5421 event
5422 }
5423
5424 fn record_transaction(
5425 connection: &Connection,
5426 project: ProjectInstanceId,
5427 instance: UsageInstanceId,
5428 owner: UsageInstanceOwner,
5429 event: &UsageEvent,
5430 policy: TelemetryRetentionPolicy,
5431 seal_after_record: bool,
5432 ) -> DbResult<()> {
5433 record_transaction_at(
5434 connection,
5435 project,
5436 instance,
5437 owner,
5438 event,
5439 policy,
5440 seal_after_record,
5441 now_epoch_seconds()?,
5442 )
5443 }
5444
5445 #[allow(clippy::too_many_arguments)]
5446 fn record_transaction_at(
5447 connection: &Connection,
5448 project: ProjectInstanceId,
5449 instance: UsageInstanceId,
5450 owner: UsageInstanceOwner,
5451 event: &UsageEvent,
5452 policy: TelemetryRetentionPolicy,
5453 seal_after_record: bool,
5454 now: i64,
5455 ) -> DbResult<()> {
5456 let transaction = Transaction::new_unchecked(connection, TransactionBehavior::Immediate)?;
5457 let result = (|| {
5458 crate::project_identity::require_bound_project_identity(&transaction, project)?;
5459 let policy = policy.validate()?;
5460 validate_event(event, policy)?;
5461 record_usage_at(
5462 &transaction,
5463 project,
5464 instance,
5465 owner,
5466 None,
5467 true,
5468 event,
5469 policy,
5470 now,
5471 seal_after_record,
5472 BaselineAdmission::BoundedRuntime,
5473 DimensionAdmission::Event,
5474 )
5475 })();
5476 match result {
5477 Ok(()) => transaction.commit().map_err(Into::into),
5478 Err(error) => {
5479 transaction.rollback()?;
5480 Err(error)
5481 }
5482 }
5483 }
5484
5485 fn scalar_count(connection: &Connection, sql: &str) -> Result<usize, Box<dyn Error>> {
5486 let value = connection.query_row(sql, [], |row| row.get::<_, i64>(0))?;
5487 Ok(usize::try_from(value)?)
5488 }
5489
5490 fn query_plan(connection: &Connection, sql: &str) -> Result<Vec<String>, Box<dyn Error>> {
5491 let mut statement = connection.prepare(&format!("EXPLAIN QUERY PLAN {sql}"))?;
5492 let rows = statement.query_map([], |row| row.get::<_, String>(3))?;
5493 Ok(rows.collect::<Result<Vec<_>, _>>()?)
5494 }
5495
5496 fn assert_plan_uses(plan: &[String], owner: &str) {
5497 assert!(
5498 plan.iter().any(|detail| detail.contains(owner)),
5499 "expected query plan to use {owner}; plan was {plan:?}"
5500 );
5501 }
5502
5503 #[test]
5504 fn retention_policy_rejects_zero_and_inverted_instance_limits() {
5505 let policy = TelemetryRetentionPolicy {
5506 prune_batch_rows: 0,
5507 ..TelemetryRetentionPolicy::default()
5508 };
5509 assert!(policy.validate().is_err());
5510
5511 let default_policy = TelemetryRetentionPolicy::default();
5512 let policy = TelemetryRetentionPolicy {
5513 max_retained_instances: default_policy.max_active_instances - 1,
5514 ..default_policy
5515 };
5516 assert!(policy.validate().is_err());
5517 }
5518
5519 #[test]
5520 fn oversized_event_fields_are_rejected_at_the_typed_storage_boundary() {
5521 type Mutator = fn(&mut UsageEvent, String);
5522
5523 let policy = TelemetryRetentionPolicy::default();
5524 let dimension_limit = policy.max_dimension_bytes;
5525 let short_witness_limit = 256.min(policy.max_baseline_witness_bytes);
5526 let cases: [(&str, usize, Mutator); 18] = [
5527 ("session_id", policy.max_label_bytes, |event, value| {
5528 event.session_id = value;
5529 }),
5530 ("command", policy.max_command_bytes, |event, value| {
5531 event.command = value;
5532 }),
5533 ("path", policy.max_path_bytes, |event, value| {
5534 event.path = Some(value);
5535 }),
5536 ("query", policy.max_query_bytes, |event, value| {
5537 event.query = Some(value);
5538 }),
5539 ("token_savings_bucket", dimension_limit, |event, value| {
5540 event.token_savings_bucket = value;
5541 }),
5542 ("provider", dimension_limit, |event, value| {
5543 event.provider = value;
5544 }),
5545 ("model", dimension_limit, |event, value| {
5546 event.model = value;
5547 }),
5548 ("tokenizer_backend", dimension_limit, |event, value| {
5549 event.tokenizer_backend = value;
5550 }),
5551 ("accuracy", dimension_limit, |event, value| {
5552 event.accuracy = value;
5553 }),
5554 ("baseline_kind", dimension_limit, |event, value| {
5555 event.baseline_kind = value;
5556 }),
5557 ("confidence", dimension_limit, |event, value| {
5558 event.confidence = value;
5559 }),
5560 ("accounting_layer", dimension_limit, |event, value| {
5561 event.accounting_layer = value;
5562 }),
5563 ("estimate_method", dimension_limit, |event, value| {
5564 event.estimate_method = value;
5565 }),
5566 ("denominator_kind", dimension_limit, |event, value| {
5567 event.denominator_kind = value;
5568 }),
5569 ("dedupe_scope", dimension_limit, |event, value| {
5570 event.dedupe_scope = value;
5571 }),
5572 ("calculation_trace", short_witness_limit, |event, value| {
5573 event.calculation_trace = value;
5574 }),
5575 (
5576 "baseline_identity",
5577 policy.max_baseline_witness_bytes,
5578 |event, value| {
5579 event.baseline_identity = value;
5580 },
5581 ),
5582 (
5583 "baseline_fingerprint",
5584 short_witness_limit,
5585 |event, value| {
5586 event.baseline_fingerprint = value;
5587 },
5588 ),
5589 ];
5590 for (field, limit, mutate) in cases {
5591 let mut candidate = event("bounded", 10, 1);
5592 mutate(&mut candidate, "x".repeat(limit + 1));
5593 assert!(matches!(
5594 validate_event(&candidate, policy),
5595 Err(DbError::TelemetryFieldTooLarge {
5596 field: found,
5597 bytes,
5598 limit: found_limit,
5599 }) if found == field && bytes == limit + 1 && found_limit == limit
5600 ));
5601 }
5602 }
5603
5604 #[test]
5605 fn spill_cleanup_is_not_applicable() {
5606 assert_eq!(
5607 SpillCleanupState::NotApplicable,
5608 SpillCleanupState::NotApplicable
5609 );
5610 assert_eq!(POLICY_VERSION, 1);
5611 assert_eq!(LOGICAL_BYTE_VERSION, 1);
5612 }
5613
5614 #[test]
5615 fn on_disk_instances_preserve_exact_aggregates_and_atomic_sealing() {
5616 let result = (|| -> Result<(), Box<dyn Error>> {
5617 let database = test_database()?;
5618 let policy = TelemetryRetentionPolicy::default();
5619 let first = instance(1)?;
5620 let second = instance(2)?;
5621 record_transaction(
5622 &database.connection,
5623 database.project,
5624 first,
5625 UsageInstanceOwner::McpProcess,
5626 &event("agent", 100, 10),
5627 policy,
5628 false,
5629 )?;
5630 record_transaction(
5631 &database.connection,
5632 database.project,
5633 first,
5634 UsageInstanceOwner::McpProcess,
5635 &event("agent", 100, 95),
5636 policy,
5637 true,
5638 )?;
5639 record_transaction(
5640 &database.connection,
5641 database.project,
5642 second,
5643 UsageInstanceOwner::McpProcess,
5644 &event("agent", 40, 10),
5645 policy,
5646 false,
5647 )?;
5648
5649 let overview =
5650 token_overview_for_project(&database.connection, database.project, Some("agent"))?;
5651 assert_eq!(overview.calls, 3);
5652 assert_eq!(overview.deduped_modeled_tokens_avoided, 25);
5653 assert_eq!(overview.repeated_baselines_deduped, 1);
5654 assert_eq!(
5655 overview.detail_availability,
5656 UsageDetailAvailability::Retained
5657 );
5658 assert_eq!(
5659 scalar_count(
5660 &database.connection,
5661 "SELECT COUNT(*) FROM usage_instances WHERE caller_label = 'agent'",
5662 )?,
5663 2
5664 );
5665 assert_eq!(
5666 scalar_count(
5667 &database.connection,
5668 "SELECT COUNT(*) FROM usage_instance_baselines",
5669 )?,
5670 1
5671 );
5672
5673 let inactive = record_transaction(
5674 &database.connection,
5675 database.project,
5676 first,
5677 UsageInstanceOwner::McpProcess,
5678 &event("agent", 20, 5),
5679 policy,
5680 false,
5681 );
5682 assert!(matches!(inactive, Err(DbError::TelemetryInstanceInactive)));
5683 let mismatched = record_transaction(
5684 &database.connection,
5685 database.project,
5686 second,
5687 UsageInstanceOwner::CliInvocation,
5688 &event("agent", 20, 5),
5689 policy,
5690 false,
5691 );
5692 assert!(matches!(
5693 mismatched,
5694 Err(DbError::TelemetryInstanceMismatch)
5695 ));
5696
5697 let database_path = database.temp.path().join("projectatlas.db");
5698 drop(database.connection);
5699 let reopened = Connection::open(database_path)?;
5700 crate::schema::initialize(&reopened, None)?;
5701 let reopened_overview =
5702 token_overview_for_project(&reopened, database.project, Some("agent"))?;
5703 assert_eq!(reopened_overview.calls, overview.calls);
5704 assert_eq!(
5705 reopened_overview.deduped_modeled_tokens_avoided,
5706 overview.deduped_modeled_tokens_avoided
5707 );
5708 Ok(())
5709 })();
5710 assert!(result.is_ok(), "on-disk telemetry test failed: {result:?}");
5711 }
5712
5713 #[test]
5714 fn sqlite_average_and_maximum_match_raw_event_accounting() {
5715 let result = (|| -> Result<(), Box<dyn Error>> {
5716 let database = test_database()?;
5717 let policy = TelemetryRetentionPolicy::default();
5718 let mut second_folder_scope = directory_event("agent", 5, 1);
5719 second_folder_scope.baseline_identity = "directory:tests".to_string();
5720 second_folder_scope.baseline_fingerprint = "directory:tests:v1".to_string();
5721 let events = vec![
5722 directory_event("agent", 101, 20),
5723 directory_event("agent", 101, 10),
5724 second_folder_scope,
5725 event("agent", 80, 20),
5726 usage_from_text(
5727 "agent",
5728 "summary",
5729 Some("src/other.rs".to_string()),
5730 None,
5731 "abcdabcd",
5732 "ab",
5733 ),
5734 ];
5735 let expected = TokenOverview::from_events(&events);
5736
5737 record_transaction(
5738 &database.connection,
5739 database.project,
5740 instance(21)?,
5741 UsageInstanceOwner::McpProcess,
5742 &events[0],
5743 policy,
5744 false,
5745 )?;
5746 record_transaction(
5747 &database.connection,
5748 database.project,
5749 instance(21)?,
5750 UsageInstanceOwner::McpProcess,
5751 &events[1],
5752 policy,
5753 true,
5754 )?;
5755 for (identity, event) in [(22, &events[2]), (23, &events[3]), (24, &events[4])] {
5756 record_transaction(
5757 &database.connection,
5758 database.project,
5759 instance(identity)?,
5760 UsageInstanceOwner::McpProcess,
5761 event,
5762 policy,
5763 true,
5764 )?;
5765 }
5766
5767 let actual =
5768 token_overview_for_project(&database.connection, database.project, Some("agent"))?;
5769 assert_eq!(actual.average_modeled_tokens_avoided, 82);
5770 assert_eq!(actual.average_tokens_avoided, 83);
5771 assert_eq!(actual.maximum_tokens_avoided, 136);
5772 assert_eq!(actual.tokens_avoided, actual.average_tokens_avoided);
5773 assert_eq!(
5774 actual.average_modeled_tokens_avoided,
5775 expected.average_modeled_tokens_avoided
5776 );
5777 assert_eq!(
5778 actual.average_tokens_avoided,
5779 expected.average_tokens_avoided
5780 );
5781 assert_eq!(
5782 actual.maximum_tokens_avoided,
5783 expected.maximum_tokens_avoided
5784 );
5785 assert_eq!(
5786 actual.deduped_modeled_tokens_avoided,
5787 expected.deduped_modeled_tokens_avoided
5788 );
5789 assert_eq!(actual.buckets, expected.buckets);
5790 Ok(())
5791 })();
5792 assert!(
5793 result.is_ok(),
5794 "SQLite average/maximum parity test failed: {result:?}"
5795 );
5796 }
5797
5798 #[test]
5799 fn directory_overflow_retains_average_policy_and_raw_parity() {
5800 let result = (|| -> Result<(), Box<dyn Error>> {
5801 let database = test_database()?;
5802 let policy = TelemetryRetentionPolicy {
5803 max_dimensions: 2,
5804 ..TelemetryRetentionPolicy::default()
5805 };
5806 let selected = event("overflow-average", 100, 10);
5807 let folder = directory_event("overflow-average", 101, 20);
5808 let expected = TokenOverview::from_events(&[selected.clone(), folder.clone()]);
5809
5810 record_transaction(
5811 &database.connection,
5812 database.project,
5813 instance(25)?,
5814 UsageInstanceOwner::McpProcess,
5815 &selected,
5816 policy,
5817 false,
5818 )?;
5819 record_transaction(
5820 &database.connection,
5821 database.project,
5822 instance(25)?,
5823 UsageInstanceOwner::McpProcess,
5824 &folder,
5825 policy,
5826 true,
5827 )?;
5828
5829 let actual = token_overview_for_project(
5830 &database.connection,
5831 database.project,
5832 Some("overflow-average"),
5833 )?;
5834 assert_eq!(actual.average_tokens_avoided, 120);
5835 assert_eq!(actual.maximum_tokens_avoided, 171);
5836 assert_eq!(
5837 actual.average_tokens_avoided,
5838 expected.average_tokens_avoided
5839 );
5840 assert_eq!(
5841 actual.maximum_tokens_avoided,
5842 expected.maximum_tokens_avoided
5843 );
5844 assert_eq!(actual.detail_availability, UsageDetailAvailability::Partial);
5845 assert_eq!(
5846 scalar_count(
5847 &database.connection,
5848 "SELECT COUNT(*) FROM usage_bucket_dimensions",
5849 )?,
5850 3
5851 );
5852 Ok(())
5853 })();
5854 assert!(
5855 result.is_ok(),
5856 "directory overflow average-policy test failed: {result:?}"
5857 );
5858 }
5859
5860 #[test]
5861 fn sqlite_and_raw_accounting_narrow_once_at_signed_bounds() {
5862 let result = (|| -> Result<(), Box<dyn Error>> {
5863 let database = test_database()?;
5864 let policy = TelemetryRetentionPolicy::default();
5865 let bound = isize::MAX as usize;
5866 let mut events = vec![
5867 event("wide", bound, 0),
5868 event("wide", bound, 0),
5869 event("wide", 0, bound),
5870 ];
5871 for (index, event) in events.iter_mut().enumerate() {
5872 event.provider = format!("wide-{index}");
5873 event.dedupe_scope = TOKEN_DEDUPE_SCOPE_EVENT.to_string();
5874 }
5875 let expected = TokenOverview::from_events(&events);
5876
5877 for (index, event) in events.iter().enumerate() {
5878 record_transaction(
5879 &database.connection,
5880 database.project,
5881 instance(30 + u8::try_from(index)?)?,
5882 UsageInstanceOwner::McpProcess,
5883 event,
5884 policy,
5885 true,
5886 )?;
5887 }
5888
5889 let actual =
5890 token_overview_for_project(&database.connection, database.project, Some("wide"))?;
5891 assert_eq!(expected.average_tokens_avoided, isize::MAX);
5892 assert_eq!(
5893 actual.average_tokens_avoided,
5894 expected.average_tokens_avoided
5895 );
5896 assert_eq!(
5897 actual.maximum_tokens_avoided,
5898 expected.maximum_tokens_avoided
5899 );
5900 assert_eq!(
5901 actual.deduped_modeled_tokens_avoided,
5902 expected.deduped_modeled_tokens_avoided
5903 );
5904 Ok(())
5905 })();
5906 assert!(
5907 result.is_ok(),
5908 "SQLite wide signed accounting parity test failed: {result:?}"
5909 );
5910 }
5911
5912 #[test]
5913 fn bounded_raw_and_label_retention_preserve_global_truth() {
5914 let result = (|| -> Result<(), Box<dyn Error>> {
5915 let database = test_database()?;
5916 let mut policy = TelemetryRetentionPolicy {
5917 max_raw_rows: 2,
5918 prune_batch_rows: 1,
5919 max_retained_labels: 1,
5920 ..TelemetryRetentionPolicy::default()
5921 };
5922 policy.checkpoint_write_interval = usize::MAX;
5923 let first = instance(3)?;
5924 record_transaction(
5925 &database.connection,
5926 database.project,
5927 first,
5928 UsageInstanceOwner::McpProcess,
5929 &event("first", 100, 10),
5930 policy,
5931 true,
5932 )?;
5933 let second = instance(4)?;
5934 record_transaction(
5935 &database.connection,
5936 database.project,
5937 second,
5938 UsageInstanceOwner::McpProcess,
5939 &event("second", 80, 20),
5940 policy,
5941 false,
5942 )?;
5943 record_transaction(
5944 &database.connection,
5945 database.project,
5946 second,
5947 UsageInstanceOwner::McpProcess,
5948 &event("second", 70, 20),
5949 policy,
5950 false,
5951 )?;
5952
5953 let state = retention_state_for_project(&database.connection, database.project)?;
5954 assert_eq!(state.raw_rows, 2);
5955 assert_eq!(state.retained_label_rows, 1);
5956 assert_eq!(state.label_tombstone_rows, 1);
5957 assert_eq!(state.spill_cleanup, SpillCleanupState::NotApplicable);
5958 let global = token_overview_for_project(&database.connection, database.project, None)?;
5959 assert_eq!(global.calls, 3);
5960 assert_eq!(global.deduped_modeled_tokens_avoided, 130);
5961 let expired =
5962 token_overview_for_project(&database.connection, database.project, Some("first"))?;
5963 assert_eq!(
5964 expired.detail_availability,
5965 UsageDetailAvailability::Expired
5966 );
5967 let unavailable = token_overview_for_project(
5968 &database.connection,
5969 database.project,
5970 Some("never-recorded"),
5971 )?;
5972 assert_eq!(
5973 unavailable.detail_availability,
5974 UsageDetailAvailability::Unavailable
5975 );
5976 let retained =
5977 token_overview_for_project(&database.connection, database.project, Some("second"))?;
5978 assert_eq!(retained.calls, 2);
5979 assert_eq!(
5980 retained.detail_availability,
5981 UsageDetailAvailability::Partial
5982 );
5983
5984 record_transaction(
5985 &database.connection,
5986 database.project,
5987 second,
5988 UsageInstanceOwner::McpProcess,
5989 &event("second", 60, 20),
5990 policy,
5991 false,
5992 )?;
5993 assert_eq!(
5994 token_overview_for_project(&database.connection, database.project, Some("second"))?
5995 .calls,
5996 3
5997 );
5998 Ok(())
5999 })();
6000 assert!(
6001 result.is_ok(),
6002 "bounded telemetry retention test failed: {result:?}"
6003 );
6004 }
6005
6006 #[test]
6007 fn raw_age_and_logical_byte_retention_preserve_exact_aggregates() {
6008 let result = (|| -> Result<(), Box<dyn Error>> {
6009 let aged = test_database()?;
6010 let aged_policy = TelemetryRetentionPolicy {
6011 max_raw_age_seconds: 10,
6012 prune_batch_rows: 1,
6013 checkpoint_write_interval: usize::MAX,
6014 ..TelemetryRetentionPolicy::default()
6015 };
6016 let aged_runtime = instance(20)?;
6017 record_transaction_at(
6018 &aged.connection,
6019 aged.project,
6020 aged_runtime,
6021 UsageInstanceOwner::McpProcess,
6022 &event("aged", 100, 10),
6023 aged_policy,
6024 false,
6025 1_000,
6026 )?;
6027 record_transaction_at(
6028 &aged.connection,
6029 aged.project,
6030 aged_runtime,
6031 UsageInstanceOwner::McpProcess,
6032 &event("aged", 100, 10),
6033 aged_policy,
6034 false,
6035 1_020,
6036 )?;
6037 let aged_state = retention_state_for_project(&aged.connection, aged.project)?;
6038 assert_eq!(aged_state.raw_rows, 1);
6039 assert_eq!(aged_state.pruned_raw_rows, 1);
6040 let aged_overview =
6041 token_overview_for_project(&aged.connection, aged.project, Some("aged"))?;
6042 assert_eq!(aged_overview.calls, 2);
6043 assert_eq!(aged_overview.deduped_modeled_tokens_avoided, 80);
6044 assert_eq!(
6045 aged_overview.detail_availability,
6046 UsageDetailAvailability::Partial
6047 );
6048
6049 let byte_bounded = test_database()?;
6050 let sample = event("bytes", 100, 10);
6051 let one_event_bytes = logical_event_bytes(&sample, Some("bytes"))?;
6052 let byte_policy = TelemetryRetentionPolicy {
6053 max_raw_logical_bytes: one_event_bytes + 1,
6054 prune_batch_rows: 1,
6055 checkpoint_write_interval: usize::MAX,
6056 ..TelemetryRetentionPolicy::default()
6057 };
6058 let byte_runtime = instance(21)?;
6059 record_transaction_at(
6060 &byte_bounded.connection,
6061 byte_bounded.project,
6062 byte_runtime,
6063 UsageInstanceOwner::McpProcess,
6064 &sample,
6065 byte_policy,
6066 false,
6067 2_000,
6068 )?;
6069 record_transaction_at(
6070 &byte_bounded.connection,
6071 byte_bounded.project,
6072 byte_runtime,
6073 UsageInstanceOwner::McpProcess,
6074 &sample,
6075 byte_policy,
6076 false,
6077 2_001,
6078 )?;
6079 let byte_state =
6080 retention_state_for_project(&byte_bounded.connection, byte_bounded.project)?;
6081 assert_eq!(byte_state.raw_rows, 1);
6082 assert!(byte_state.raw_logical_bytes <= byte_policy.max_raw_logical_bytes);
6083 assert_eq!(byte_state.pruned_raw_rows, 1);
6084 let byte_overview = token_overview_for_project(
6085 &byte_bounded.connection,
6086 byte_bounded.project,
6087 Some("bytes"),
6088 )?;
6089 assert_eq!(byte_overview.calls, 2);
6090 assert_eq!(byte_overview.deduped_modeled_tokens_avoided, 80);
6091 Ok(())
6092 })();
6093 assert!(
6094 result.is_ok(),
6095 "raw telemetry budget test failed: {result:?}"
6096 );
6097 }
6098
6099 #[test]
6100 fn production_batch_raw_retention_removes_only_the_required_oldest_prefix() {
6101 let result = (|| -> Result<(), Box<dyn Error>> {
6102 let aged = test_database()?;
6103 let aged_policy = TelemetryRetentionPolicy {
6104 max_raw_rows: 10,
6105 max_raw_age_seconds: 10,
6106 prune_batch_rows: TelemetryRetentionPolicy::default().prune_batch_rows,
6107 checkpoint_write_interval: usize::MAX,
6108 ..TelemetryRetentionPolicy::default()
6109 };
6110 let runtime = instance(22)?;
6111 for now in [1_000, 1_020, 1_021] {
6112 record_transaction_at(
6113 &aged.connection,
6114 aged.project,
6115 runtime,
6116 UsageInstanceOwner::McpProcess,
6117 &event("aged-prefix", 100, 10),
6118 aged_policy,
6119 false,
6120 now,
6121 )?;
6122 }
6123 let epochs = aged
6124 .connection
6125 .prepare("SELECT created_at_epoch FROM usage_events ORDER BY created_at_epoch")?
6126 .query_map([], |row| row.get::<_, i64>(0))?
6127 .collect::<Result<Vec<_>, _>>()?;
6128 assert_eq!(epochs, vec![1_020, 1_021]);
6129
6130 let capped = test_database()?;
6131 let capped_policy = TelemetryRetentionPolicy {
6132 max_raw_rows: 2,
6133 prune_batch_rows: TelemetryRetentionPolicy::default().prune_batch_rows,
6134 checkpoint_write_interval: usize::MAX,
6135 ..TelemetryRetentionPolicy::default()
6136 };
6137 let runtime = instance(23)?;
6138 for now in [2_000, 2_001, 2_002] {
6139 record_transaction_at(
6140 &capped.connection,
6141 capped.project,
6142 runtime,
6143 UsageInstanceOwner::McpProcess,
6144 &event("row-prefix", 100, 10),
6145 capped_policy,
6146 false,
6147 now,
6148 )?;
6149 }
6150 let epochs = capped
6151 .connection
6152 .prepare("SELECT created_at_epoch FROM usage_events ORDER BY created_at_epoch")?
6153 .query_map([], |row| row.get::<_, i64>(0))?
6154 .collect::<Result<Vec<_>, _>>()?;
6155 assert_eq!(epochs, vec![2_001, 2_002]);
6156 assert_eq!(
6157 retention_state_for_project(&capped.connection, capped.project)?.pruned_raw_rows,
6158 1
6159 );
6160 Ok(())
6161 })();
6162 assert!(
6163 result.is_ok(),
6164 "exact raw-prefix retention test failed: {result:?}"
6165 );
6166 }
6167
6168 #[test]
6169 fn production_batch_tombstone_retention_removes_only_exact_excess() {
6170 let result = (|| -> Result<(), Box<dyn Error>> {
6171 let database = test_database()?;
6172 let policy = TelemetryRetentionPolicy {
6173 max_label_tombstones: 2,
6174 max_instance_tombstones: 2,
6175 prune_batch_rows: TelemetryRetentionPolicy::default().prune_batch_rows,
6176 checkpoint_write_interval: usize::MAX,
6177 ..TelemetryRetentionPolicy::default()
6178 };
6179 for (offset, label) in ["first", "second", "third"].into_iter().enumerate() {
6180 upsert_label_tombstone(
6181 &database.connection,
6182 database.project.as_bytes().as_slice(),
6183 label,
6184 1_000 + i64::try_from(offset)?,
6185 None,
6186 )?;
6187 }
6188 for (offset, runtime) in [31_u8, 32, 33].into_iter().enumerate() {
6189 database.connection.execute(
6190 "INSERT INTO usage_instance_tombstones(
6191 project_instance_id, runtime_instance_id, retired_at_epoch
6192 ) VALUES(?1, ?2, ?3)",
6193 params![
6194 database.project.as_bytes().as_slice(),
6195 instance(runtime)?.as_bytes().as_slice(),
6196 1_000 + i64::try_from(offset)?,
6197 ],
6198 )?;
6199 }
6200 increment_retention_counter(
6201 &database.connection,
6202 RetentionCounter::InstanceTombstoneRows,
6203 3,
6204 )?;
6205 let deleted = prune_tombstones_once(&database.connection, policy, 1_100)?;
6206 assert_eq!(deleted, 2);
6207 assert_eq!(
6208 scalar_count(
6209 &database.connection,
6210 "SELECT COUNT(*) FROM usage_label_tombstones",
6211 )?,
6212 2
6213 );
6214 assert_eq!(
6215 scalar_count(
6216 &database.connection,
6217 "SELECT COUNT(*) FROM usage_instance_tombstones",
6218 )?,
6219 2
6220 );
6221 assert_eq!(
6222 database.connection.query_row(
6223 "SELECT caller_label FROM usage_label_tombstones
6224 ORDER BY expired_at_epoch LIMIT 1",
6225 [],
6226 |row| row.get::<_, String>(0),
6227 )?,
6228 "second"
6229 );
6230 Ok(())
6231 })();
6232 assert!(
6233 result.is_ok(),
6234 "exact tombstone retention test failed: {result:?}"
6235 );
6236 }
6237
6238 #[test]
6239 fn daily_capacity_prunes_old_history_before_reserving_both_current_rows() {
6240 let result = (|| -> Result<(), Box<dyn Error>> {
6241 let database = test_database()?;
6242 let policy = TelemetryRetentionPolicy {
6243 max_daily_rows: 4,
6244 retained_trend_days: 1,
6245 prune_batch_rows: TelemetryRetentionPolicy::default().prune_batch_rows,
6246 checkpoint_write_interval: usize::MAX,
6247 ..TelemetryRetentionPolicy::default()
6248 };
6249 let runtime = instance(24)?;
6250 let mut first = event("daily", 100, 10);
6251 first.provider = "first-provider".to_string();
6252 let mut second = event("daily", 90, 10);
6253 second.provider = "second-provider".to_string();
6254 for value in [&first, &second] {
6255 record_transaction_at(
6256 &database.connection,
6257 database.project,
6258 runtime,
6259 UsageInstanceOwner::McpProcess,
6260 value,
6261 policy,
6262 false,
6263 SECONDS_PER_DAY,
6264 )?;
6265 }
6266 assert_eq!(
6267 retention_counter(&database.connection, RetentionCounter::DailyRows)?,
6268 4
6269 );
6270
6271 let mut current = event("daily", 80, 10);
6272 current.provider = "current-provider".to_string();
6273 let current_epoch = 3 * SECONDS_PER_DAY;
6274 record_transaction_at(
6275 &database.connection,
6276 database.project,
6277 runtime,
6278 UsageInstanceOwner::McpProcess,
6279 ¤t,
6280 policy,
6281 false,
6282 current_epoch,
6283 )?;
6284 assert_eq!(
6285 retention_counter(&database.connection, RetentionCounter::DailyRows)?,
6286 2
6287 );
6288 let current_rows = database.connection.query_row(
6289 "SELECT
6290 (SELECT COUNT(*) FROM usage_daily_aggregates AS aggregate
6291 JOIN usage_bucket_dimensions AS dimension USING(dimension_id)
6292 WHERE aggregate.day_epoch = ?1 AND dimension.provider = ?2),
6293 (SELECT COUNT(*) FROM usage_instance_daily_aggregates AS aggregate
6294 JOIN usage_bucket_dimensions AS dimension USING(dimension_id)
6295 WHERE aggregate.day_epoch = ?1 AND dimension.provider = ?2)",
6296 params![current_epoch, "current-provider"],
6297 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
6298 )?;
6299 assert_eq!(current_rows, (1, 1));
6300 let trends = token_trends_for_project(
6301 &database.connection,
6302 database.project,
6303 Some("daily"),
6304 TokenTrendWindow::Day,
6305 )?;
6306 assert_eq!(trends.periods.len(), 1);
6307 assert_eq!(trends.periods[0].calls, 1);
6308 Ok(())
6309 })();
6310 assert!(
6311 result.is_ok(),
6312 "daily capacity reservation test failed: {result:?}"
6313 );
6314 }
6315
6316 #[test]
6317 fn label_reports_group_many_instances_inside_sqlite() {
6318 let result = (|| -> Result<(), Box<dyn Error>> {
6319 let database = test_database()?;
6320 let policy = TelemetryRetentionPolicy {
6321 checkpoint_write_interval: usize::MAX,
6322 ..TelemetryRetentionPolicy::default()
6323 };
6324 for runtime in 1_u8..=48 {
6325 record_transaction_at(
6326 &database.connection,
6327 database.project,
6328 instance(runtime)?,
6329 UsageInstanceOwner::McpProcess,
6330 &event("grouped", 100, 10),
6331 policy,
6332 true,
6333 10 * SECONDS_PER_DAY,
6334 )?;
6335 }
6336 let overview = token_overview_for_project(
6337 &database.connection,
6338 database.project,
6339 Some("grouped"),
6340 )?;
6341 assert_eq!(overview.calls, 48);
6342 assert_eq!(overview.buckets.len(), 1);
6343 let trends = token_trends_for_project(
6344 &database.connection,
6345 database.project,
6346 Some("grouped"),
6347 TokenTrendWindow::Day,
6348 )?;
6349 assert_eq!(trends.periods.len(), 1);
6350 assert_eq!(trends.periods[0].calls, 48);
6351 assert_eq!(trends.periods[0].buckets.len(), 1);
6352 Ok(())
6353 })();
6354 assert!(
6355 result.is_ok(),
6356 "SQLite-side telemetry grouping test failed: {result:?}"
6357 );
6358 }
6359
6360 #[test]
6361 fn new_project_runtime_reclaims_inactive_instance_and_label_capacity() {
6362 let result = (|| -> Result<(), Box<dyn Error>> {
6363 let database = test_database()?;
6364 let policy = TelemetryRetentionPolicy {
6365 max_active_instances: 1,
6366 max_retained_instances: 1,
6367 max_retained_labels: 1,
6368 ..TelemetryRetentionPolicy::default()
6369 };
6370 let old_project = database.project;
6371 let old_runtime = instance(30)?;
6372 record_transaction(
6373 &database.connection,
6374 old_project,
6375 old_runtime,
6376 UsageInstanceOwner::McpProcess,
6377 &event("old-project", 100, 10),
6378 policy,
6379 false,
6380 )?;
6381 assert_eq!(
6382 seal_project_usage_instances(&database.connection, old_project)?,
6383 1
6384 );
6385 let sealed_state = retention_state_for_project(&database.connection, old_project)?;
6386 assert_eq!(sealed_state.baseline_rows, 0);
6387
6388 let new_project = ProjectInstanceId::from_bytes([31; 16])?;
6389 database.connection.execute(
6390 "UPDATE project_identity SET project_instance_id = ?1 WHERE singleton = 1",
6391 [new_project.as_bytes().as_slice()],
6392 )?;
6393 let new_runtime = old_runtime;
6394 record_transaction(
6395 &database.connection,
6396 new_project,
6397 new_runtime,
6398 UsageInstanceOwner::McpProcess,
6399 &event("new-project", 80, 20),
6400 policy,
6401 false,
6402 )?;
6403
6404 assert_eq!(
6405 scalar_count(&database.connection, "SELECT COUNT(*) FROM usage_instances")?,
6406 1
6407 );
6408 assert_eq!(
6409 scalar_count(&database.connection, "SELECT COUNT(*) FROM usage_labels")?,
6410 1
6411 );
6412 assert_eq!(
6413 scalar_count(
6414 &database.connection,
6415 "SELECT COUNT(*) FROM usage_global_aggregates",
6416 )?,
6417 1
6418 );
6419 assert_eq!(
6420 scalar_count(
6421 &database.connection,
6422 "SELECT COUNT(*) FROM usage_labels WHERE caller_label = 'new-project'",
6423 )?,
6424 1
6425 );
6426 let old_instance_tombstones = database.connection.query_row(
6427 "SELECT COUNT(*) FROM usage_instance_tombstones
6428 WHERE project_instance_id = ?1",
6429 [old_project.as_bytes().as_slice()],
6430 |row| row.get::<_, i64>(0),
6431 )?;
6432 assert_eq!(old_instance_tombstones, 1);
6433 let old_label_tombstones = database.connection.query_row(
6434 "SELECT COUNT(*) FROM usage_label_tombstones
6435 WHERE project_instance_id = ?1",
6436 [old_project.as_bytes().as_slice()],
6437 |row| row.get::<_, i64>(0),
6438 )?;
6439 assert_eq!(old_label_tombstones, 1);
6440 let overview = token_overview_for_project(&database.connection, new_project, None)?;
6441 assert_eq!(overview.calls, 1);
6442 Ok(())
6443 })();
6444 assert!(
6445 result.is_ok(),
6446 "project-rotation telemetry capacity test failed: {result:?}"
6447 );
6448 }
6449
6450 #[test]
6451 fn active_labels_and_reserved_overflow_dimension_remain_disjoint() {
6452 let result = (|| -> Result<(), Box<dyn Error>> {
6453 let database = test_database()?;
6454 let policy = TelemetryRetentionPolicy {
6455 max_retained_labels: 1,
6456 max_dimensions: 2,
6457 ..TelemetryRetentionPolicy::default()
6458 };
6459 let active = instance(5)?;
6460 let mut reserved = event("active", 100, 10);
6461 reserved.token_savings_bucket = OVERFLOW_DIMENSION.to_string();
6462 reserved.provider = OVERFLOW_DIMENSION.to_string();
6463 reserved.model = OVERFLOW_DIMENSION.to_string();
6464 reserved.tokenizer_backend = OVERFLOW_DIMENSION.to_string();
6465 reserved.accuracy = OVERFLOW_DIMENSION.to_string();
6466 reserved.baseline_kind = OVERFLOW_DIMENSION.to_string();
6467 reserved.confidence = OVERFLOW_DIMENSION.to_string();
6468 reserved.accounting_layer = OVERFLOW_DIMENSION.to_string();
6469 reserved.estimate_method = OVERFLOW_DIMENSION.to_string();
6470 reserved.denominator_kind = OVERFLOW_DIMENSION.to_string();
6471 reserved.dedupe_scope = OVERFLOW_DIMENSION.to_string();
6472 record_transaction(
6473 &database.connection,
6474 database.project,
6475 active,
6476 UsageInstanceOwner::McpProcess,
6477 &reserved,
6478 policy,
6479 false,
6480 )?;
6481 assert_eq!(
6482 scalar_count(
6483 &database.connection,
6484 "SELECT COUNT(*) FROM usage_bucket_dimensions
6485 WHERE token_savings_bucket = '<overflow>'",
6486 )?,
6487 2
6488 );
6489 assert_eq!(
6490 scalar_count(
6491 &database.connection,
6492 "SELECT COUNT(DISTINCT overflow) FROM usage_bucket_dimensions
6493 WHERE token_savings_bucket = '<overflow>'",
6494 )?,
6495 2
6496 );
6497
6498 let rejected = record_transaction(
6499 &database.connection,
6500 database.project,
6501 instance(6)?,
6502 UsageInstanceOwner::McpProcess,
6503 &event("replacement", 50, 10),
6504 policy,
6505 false,
6506 );
6507 assert!(matches!(rejected, Err(DbError::TelemetryInstanceCapacity)));
6508 assert_eq!(
6509 scalar_count(
6510 &database.connection,
6511 "SELECT COUNT(*) FROM usage_labels WHERE caller_label = 'active'",
6512 )?,
6513 1
6514 );
6515 assert_eq!(
6516 scalar_count(
6517 &database.connection,
6518 "SELECT COUNT(*) FROM usage_label_tombstones",
6519 )?,
6520 0
6521 );
6522 Ok(())
6523 })();
6524 assert!(
6525 result.is_ok(),
6526 "active-label overflow sentinel test failed: {result:?}"
6527 );
6528 }
6529
6530 #[test]
6531 fn dimension_capacity_uses_reserved_overflow_and_reports_partial_detail() {
6532 let result = (|| -> Result<(), Box<dyn Error>> {
6533 let database = test_database()?;
6534 let policy = TelemetryRetentionPolicy {
6535 max_dimensions: 2,
6536 ..TelemetryRetentionPolicy::default()
6537 };
6538 let runtime = instance(35)?;
6539 record_transaction(
6540 &database.connection,
6541 database.project,
6542 runtime,
6543 UsageInstanceOwner::McpProcess,
6544 &event("dimensions", 100, 10),
6545 policy,
6546 false,
6547 )?;
6548 let mut overflowed = event("dimensions", 50, 10);
6549 overflowed.provider = "another-provider".to_string();
6550 overflowed.baseline_identity = "source:src/other.rs".to_string();
6551 overflowed.baseline_fingerprint = "source:src/other.rs:v1".to_string();
6552 record_transaction(
6553 &database.connection,
6554 database.project,
6555 runtime,
6556 UsageInstanceOwner::McpProcess,
6557 &overflowed,
6558 policy,
6559 false,
6560 )?;
6561
6562 assert_eq!(
6563 scalar_count(
6564 &database.connection,
6565 "SELECT COUNT(*) FROM usage_bucket_dimensions",
6566 )?,
6567 2
6568 );
6569 let overflow_calls = database.connection.query_row(
6570 "SELECT aggregate.calls
6571 FROM usage_global_aggregates AS aggregate
6572 JOIN usage_bucket_dimensions AS dimension USING(dimension_id)
6573 WHERE aggregate.project_instance_id = ?1 AND dimension.overflow = 1",
6574 [database.project.as_bytes().as_slice()],
6575 |row| row.get::<_, i64>(0),
6576 )?;
6577 assert_eq!(overflow_calls, 1);
6578 let overview = token_overview_for_project(
6579 &database.connection,
6580 database.project,
6581 Some("dimensions"),
6582 )?;
6583 assert_eq!(overview.calls, 2);
6584 assert_eq!(
6585 overview.detail_availability,
6586 UsageDetailAvailability::Partial
6587 );
6588 assert_eq!(
6589 overview.average_policy.evidence,
6590 TOKEN_AVERAGE_POLICY_OVERFLOW_EVIDENCE
6591 );
6592 Ok(())
6593 })();
6594 assert!(
6595 result.is_ok(),
6596 "dimension overflow retention test failed: {result:?}"
6597 );
6598 }
6599
6600 #[test]
6601 fn rejected_events_and_read_only_reports_leave_storage_unchanged() {
6602 let result = (|| -> Result<(), Box<dyn Error>> {
6603 let database = test_database()?;
6604 let policy = TelemetryRetentionPolicy::default();
6605 let runtime = instance(7)?;
6606 record_transaction(
6607 &database.connection,
6608 database.project,
6609 runtime,
6610 UsageInstanceOwner::McpProcess,
6611 &event("reader", 100, 10),
6612 policy,
6613 false,
6614 )?;
6615 let counters_before = database.connection.query_row(
6616 "SELECT raw_rows, baseline_rows, instance_rows, daily_rows,
6617 writes_since_checkpoint
6618 FROM usage_retention_state WHERE singleton = 1",
6619 [],
6620 |row| {
6621 Ok((
6622 row.get::<_, i64>(0)?,
6623 row.get::<_, i64>(1)?,
6624 row.get::<_, i64>(2)?,
6625 row.get::<_, i64>(3)?,
6626 row.get::<_, i64>(4)?,
6627 ))
6628 },
6629 )?;
6630 let mut oversized = event("reader", 10, 5);
6631 oversized.query = Some("x".repeat(policy.max_query_bytes + 1));
6632 let rejected = record_transaction(
6633 &database.connection,
6634 database.project,
6635 runtime,
6636 UsageInstanceOwner::McpProcess,
6637 &oversized,
6638 policy,
6639 false,
6640 );
6641 assert!(matches!(
6642 rejected,
6643 Err(DbError::TelemetryFieldTooLarge { .. })
6644 ));
6645 let _ =
6646 usage_events_for_project(&database.connection, database.project, Some("reader"))?;
6647 let _ =
6648 token_overview_for_project(&database.connection, database.project, Some("reader"))?;
6649 let _ = token_trends_for_project(
6650 &database.connection,
6651 database.project,
6652 Some("reader"),
6653 TokenTrendWindow::Day,
6654 )?;
6655 let state = retention_state_for_project(&database.connection, database.project)?;
6656 assert_eq!(state.spill_cleanup, SpillCleanupState::NotApplicable);
6657 let counters_after = database.connection.query_row(
6658 "SELECT raw_rows, baseline_rows, instance_rows, daily_rows,
6659 writes_since_checkpoint
6660 FROM usage_retention_state WHERE singleton = 1",
6661 [],
6662 |row| {
6663 Ok((
6664 row.get::<_, i64>(0)?,
6665 row.get::<_, i64>(1)?,
6666 row.get::<_, i64>(2)?,
6667 row.get::<_, i64>(3)?,
6668 row.get::<_, i64>(4)?,
6669 ))
6670 },
6671 )?;
6672 assert_eq!(counters_after, counters_before);
6673
6674 let storage_before = database.connection.query_row(
6675 "SELECT
6676 (SELECT COUNT(*) FROM usage_events),
6677 (SELECT COUNT(*) FROM usage_instance_baselines),
6678 (SELECT calls FROM usage_global_aggregates),
6679 raw_rows,
6680 baseline_rows
6681 FROM usage_retention_state WHERE singleton = 1",
6682 [],
6683 |row| {
6684 Ok((
6685 row.get::<_, i64>(0)?,
6686 row.get::<_, i64>(1)?,
6687 row.get::<_, i64>(2)?,
6688 row.get::<_, i64>(3)?,
6689 row.get::<_, i64>(4)?,
6690 ))
6691 },
6692 )?;
6693 database.connection.execute_batch(
6694 "CREATE TEMP TRIGGER reject_telemetry_aggregate_update
6695 BEFORE UPDATE ON usage_global_aggregates
6696 BEGIN
6697 SELECT RAISE(ABORT, 'injected telemetry aggregate failure');
6698 END;",
6699 )?;
6700 let mut late_failure = event("reader", 200, 20);
6701 late_failure.baseline_identity = "source:src/lib.rs:second".to_string();
6702 late_failure.baseline_fingerprint = "source:src/lib.rs:v2".to_string();
6703 assert!(
6704 record_transaction(
6705 &database.connection,
6706 database.project,
6707 runtime,
6708 UsageInstanceOwner::McpProcess,
6709 &late_failure,
6710 policy,
6711 false,
6712 )
6713 .is_err()
6714 );
6715 database
6716 .connection
6717 .execute_batch("DROP TRIGGER reject_telemetry_aggregate_update")?;
6718 let storage_after = database.connection.query_row(
6719 "SELECT
6720 (SELECT COUNT(*) FROM usage_events),
6721 (SELECT COUNT(*) FROM usage_instance_baselines),
6722 (SELECT calls FROM usage_global_aggregates),
6723 raw_rows,
6724 baseline_rows
6725 FROM usage_retention_state WHERE singleton = 1",
6726 [],
6727 |row| {
6728 Ok((
6729 row.get::<_, i64>(0)?,
6730 row.get::<_, i64>(1)?,
6731 row.get::<_, i64>(2)?,
6732 row.get::<_, i64>(3)?,
6733 row.get::<_, i64>(4)?,
6734 ))
6735 },
6736 )?;
6737 assert_eq!(storage_after, storage_before);
6738
6739 let lookalike = database.temp.path().join("usage-telemetry-spill.db");
6740 std::fs::write(&lookalike, b"not owned by ProjectAtlas")?;
6741 maintain_after_commit_for_project(
6742 &database.connection,
6743 None,
6744 database.project,
6745 policy,
6746 )?;
6747 assert_eq!(std::fs::read(lookalike)?, b"not owned by ProjectAtlas");
6748 Ok(())
6749 })();
6750 assert!(
6751 result.is_ok(),
6752 "telemetry rejection rollback test failed: {result:?}"
6753 );
6754 }
6755
6756 #[test]
6757 fn public_raw_event_read_rejects_negative_optional_token_estimate() -> Result<(), Box<dyn Error>>
6758 {
6759 let database = production_database()?;
6760 database
6761 .store
6762 .record_usage(&event("corrupt-reader", 100, 10))?;
6763 database.store.connection.execute(
6764 "UPDATE usage_events
6765 SET estimated_tokens_without_projectatlas = -1
6766 WHERE id = (SELECT id FROM usage_events ORDER BY id DESC LIMIT 1)",
6767 [],
6768 )?;
6769
6770 let Err(error) = database.store.usage_events(Some("corrupt-reader")) else {
6771 return Err(io::Error::other("negative persisted token estimate was accepted").into());
6772 };
6773 require(
6774 matches!(
6775 error,
6776 DbError::TelemetryIntegerOverflow {
6777 field: "estimated_tokens_without_projectatlas"
6778 }
6779 ),
6780 "negative persisted token estimate returned the wrong error",
6781 )?;
6782 Ok(())
6783 }
6784
6785 #[test]
6786 fn baseline_capacity_collision_and_integer_overflow_roll_back_completely() {
6787 let result = (|| -> Result<(), Box<dyn Error>> {
6788 let database = test_database()?;
6789 let policy = TelemetryRetentionPolicy {
6790 max_baselines_per_instance: 1,
6791 ..TelemetryRetentionPolicy::default()
6792 };
6793
6794 let capacity_runtime = instance(40)?;
6795 record_transaction(
6796 &database.connection,
6797 database.project,
6798 capacity_runtime,
6799 UsageInstanceOwner::McpProcess,
6800 &event("capacity", 100, 10),
6801 policy,
6802 false,
6803 )?;
6804 let capacity_before = database.connection.query_row(
6805 "SELECT
6806 (SELECT COUNT(*) FROM usage_events),
6807 (SELECT COUNT(*) FROM usage_instance_baselines),
6808 (SELECT SUM(calls) FROM usage_global_aggregates),
6809 raw_rows,
6810 baseline_rows
6811 FROM usage_retention_state WHERE singleton = 1",
6812 [],
6813 |row| {
6814 Ok((
6815 row.get::<_, i64>(0)?,
6816 row.get::<_, i64>(1)?,
6817 row.get::<_, i64>(2)?,
6818 row.get::<_, i64>(3)?,
6819 row.get::<_, i64>(4)?,
6820 ))
6821 },
6822 )?;
6823 let mut second_baseline = event("capacity", 90, 10);
6824 second_baseline.baseline_identity = "source:src/other.rs".to_string();
6825 second_baseline.baseline_fingerprint = "source:src/other.rs:v1".to_string();
6826 assert!(matches!(
6827 record_transaction(
6828 &database.connection,
6829 database.project,
6830 capacity_runtime,
6831 UsageInstanceOwner::McpProcess,
6832 &second_baseline,
6833 policy,
6834 false,
6835 ),
6836 Err(DbError::TelemetryBaselineCapacity)
6837 ));
6838 let capacity_after = database.connection.query_row(
6839 "SELECT
6840 (SELECT COUNT(*) FROM usage_events),
6841 (SELECT COUNT(*) FROM usage_instance_baselines),
6842 (SELECT SUM(calls) FROM usage_global_aggregates),
6843 raw_rows,
6844 baseline_rows
6845 FROM usage_retention_state WHERE singleton = 1",
6846 [],
6847 |row| {
6848 Ok((
6849 row.get::<_, i64>(0)?,
6850 row.get::<_, i64>(1)?,
6851 row.get::<_, i64>(2)?,
6852 row.get::<_, i64>(3)?,
6853 row.get::<_, i64>(4)?,
6854 ))
6855 },
6856 )?;
6857 assert_eq!(capacity_after, capacity_before);
6858
6859 let collision_runtime = instance(41)?;
6860 let collision_event = event("collision", 60, 10);
6861 record_transaction(
6862 &database.connection,
6863 database.project,
6864 collision_runtime,
6865 UsageInstanceOwner::McpProcess,
6866 &collision_event,
6867 TelemetryRetentionPolicy::default(),
6868 false,
6869 )?;
6870 database.connection.execute(
6871 "UPDATE usage_instance_baselines
6872 SET baseline_identity = 'source:src/foo.rs'
6873 WHERE instance_row_id = (
6874 SELECT instance_row_id FROM usage_instances
6875 WHERE project_instance_id = ?1 AND runtime_instance_id = ?2
6876 )",
6877 params![
6878 database.project.as_bytes().as_slice(),
6879 collision_runtime.as_bytes().as_slice()
6880 ],
6881 )?;
6882 let collision_raw_before =
6883 scalar_count(&database.connection, "SELECT COUNT(*) FROM usage_events")?;
6884 assert!(matches!(
6885 record_transaction(
6886 &database.connection,
6887 database.project,
6888 collision_runtime,
6889 UsageInstanceOwner::McpProcess,
6890 &collision_event,
6891 TelemetryRetentionPolicy::default(),
6892 false,
6893 ),
6894 Err(DbError::TelemetryBaselineCollision)
6895 ));
6896 assert_eq!(
6897 scalar_count(&database.connection, "SELECT COUNT(*) FROM usage_events")?,
6898 collision_raw_before
6899 );
6900
6901 let overflow_runtime = instance(42)?;
6902 record_transaction(
6903 &database.connection,
6904 database.project,
6905 overflow_runtime,
6906 UsageInstanceOwner::McpProcess,
6907 &event("overflow", 50, 10),
6908 TelemetryRetentionPolicy::default(),
6909 false,
6910 )?;
6911 database
6912 .connection
6913 .execute("UPDATE usage_global_aggregates SET calls = ?1", [i64::MAX])?;
6914 let overflow_raw_before =
6915 scalar_count(&database.connection, "SELECT COUNT(*) FROM usage_events")?;
6916 let overflow_baselines_before = scalar_count(
6917 &database.connection,
6918 "SELECT COUNT(*) FROM usage_instance_baselines",
6919 )?;
6920 assert!(matches!(
6921 record_transaction(
6922 &database.connection,
6923 database.project,
6924 overflow_runtime,
6925 UsageInstanceOwner::McpProcess,
6926 &event("overflow", 50, 10),
6927 TelemetryRetentionPolicy::default(),
6928 false,
6929 ),
6930 Err(DbError::TelemetryIntegerOverflow {
6931 field: AGGREGATE_COUNTER_FIELD
6932 })
6933 ));
6934 assert_eq!(
6935 scalar_count(&database.connection, "SELECT COUNT(*) FROM usage_events")?,
6936 overflow_raw_before
6937 );
6938 assert_eq!(
6939 scalar_count(
6940 &database.connection,
6941 "SELECT COUNT(*) FROM usage_instance_baselines",
6942 )?,
6943 overflow_baselines_before
6944 );
6945 Ok(())
6946 })();
6947 assert!(
6948 result.is_ok(),
6949 "telemetry baseline rollback test failed: {result:?}"
6950 );
6951 }
6952
6953 #[test]
6954 fn production_store_reports_live_page_and_connection_policy_state() {
6955 let result = (|| -> Result<(), Box<dyn Error>> {
6956 let mut database = production_database()?;
6957 let initial = database.store.telemetry_retention_state()?;
6958 assert!(initial.page_count > 0);
6959 assert_eq!(initial.checkpoint_state, TelemetryCheckpointState::NotDue);
6960 assert_eq!(
6961 initial.statistics_policy,
6962 PlannerStatisticsPolicy::NotConfigured
6963 );
6964 assert_eq!(
6965 initial.statistics_state,
6966 PlannerStatisticsState::NotInitialized
6967 );
6968 assert_eq!(
6969 initial.connection_busy_timeout_ms,
6970 initial.normal_busy_timeout_ms
6971 );
6972 assert_eq!(initial.normal_busy_timeout_ms, 5_000);
6973 assert_eq!(initial.telemetry_busy_timeout_ms, 25);
6974 assert!(initial.wal_autocheckpoint_pages > 0);
6975
6976 let nodes = (0..1_024)
6977 .map(|index| {
6978 let path = format!("src/generated/{index:04}.rs");
6979 Node {
6980 parent_path: normalized_parent(&path),
6981 path,
6982 kind: NodeKind::File,
6983 extension: Some(".rs".to_string()),
6984 language: Some("rust".to_string()),
6985 size_bytes: Some(4_096),
6986 mtime_ns: Some(i64::from(index)),
6987 content_hash: Some(format!("hash-{index:04}")),
6988 }
6989 })
6990 .collect::<Vec<_>>();
6991 database.store.replace_scan(&nodes)?;
6992 let grown = database.store.telemetry_retention_state()?;
6993 let live_grown_pages = pragma_count(&database.store.connection, "page_count")?;
6994 assert_eq!(
6995 grown.page_count,
6996 count_usize("page_count", live_grown_pages)?
6997 );
6998 assert!(grown.page_count >= initial.page_count);
6999
7000 database.store.replace_scan(&[])?;
7001 let deleted = database.store.telemetry_retention_state()?;
7002 assert_eq!(
7003 deleted.freelist_pages,
7004 count_usize(
7005 "freelist_pages",
7006 pragma_count(&database.store.connection, "freelist_count")?,
7007 )?
7008 );
7009 assert_eq!(
7010 deleted.page_count,
7011 count_usize(
7012 "page_count",
7013 pragma_count(&database.store.connection, "page_count")?,
7014 )?
7015 );
7016 Ok(())
7017 })();
7018 assert!(
7019 result.is_ok(),
7020 "production page-policy report test failed: {result:?}"
7021 );
7022 }
7023
7024 #[test]
7025 fn checkpoint_finalization_preserves_writes_committed_after_attempt() {
7026 assert_eq!(
7027 writes_after_checkpoint_attempt(
7028 TelemetryCheckpointState::Completed,
7029 1_024,
7030 1_025,
7031 true,
7032 ),
7033 1
7034 );
7035 assert_eq!(
7036 writes_after_checkpoint_attempt(
7037 TelemetryCheckpointState::Completed,
7038 1_024,
7039 1_025,
7040 false,
7041 ),
7042 1_025
7043 );
7044 assert_eq!(
7045 writes_after_checkpoint_attempt(TelemetryCheckpointState::Busy, 1_024, 1_025, true,),
7046 1_025
7047 );
7048 assert_eq!(
7049 writes_after_checkpoint_attempt(TelemetryCheckpointState::Error, 1_024, 1_025, true,),
7050 1_025
7051 );
7052 }
7053
7054 #[test]
7055 fn overlapping_checkpoint_finalizers_preserve_later_event_debt() {
7056 let result = (|| -> Result<(), Box<dyn Error>> {
7057 let database = production_database()?;
7058 let policy = TelemetryRetentionPolicy::default();
7059 let first = AtlasStore::open_for_project(&database.database_path, &database.root)?;
7060 let second = AtlasStore::open_for_project(&database.database_path, &database.root)?;
7061 let event_writer =
7062 AtlasStore::open_for_project(&database.database_path, &database.root)?;
7063 let first_instance = instance(91)?;
7064 let second_instance = instance(92)?;
7065 let first_event = event("checkpoint-overlap-first", 100, 10);
7066 let second_event = event("checkpoint-overlap-second", 100, 10);
7067 database.store.connection.execute(
7068 "UPDATE usage_retention_state
7069 SET writes_since_checkpoint = ?1 WHERE singleton = 1",
7070 [to_i64(
7071 "checkpoint_write_interval",
7072 policy.checkpoint_write_interval,
7073 )?],
7074 )?;
7075
7076 maintain_after_commit_for_project_with_checkpoint(
7077 &first.connection,
7078 first.validated_project_root.as_deref(),
7079 database.project,
7080 policy,
7081 |connection| {
7082 let state = passive_checkpoint_state(connection);
7083 record_transaction(
7084 &event_writer.connection,
7085 database.project,
7086 first_instance,
7087 UsageInstanceOwner::McpProcess,
7088 &first_event,
7089 policy,
7090 false,
7091 )?;
7092 maintain_after_commit_for_project_with_checkpoint(
7093 &second.connection,
7094 second.validated_project_root.as_deref(),
7095 database.project,
7096 policy,
7097 |connection| {
7098 let state = passive_checkpoint_state(connection);
7099 record_transaction(
7100 &event_writer.connection,
7101 database.project,
7102 second_instance,
7103 UsageInstanceOwner::McpProcess,
7104 &second_event,
7105 policy,
7106 false,
7107 )?;
7108 Ok(state)
7109 },
7110 )?;
7111 Ok(state)
7112 },
7113 )?;
7114
7115 let retention = database.store.telemetry_retention_state()?;
7116 assert_eq!(
7117 retention.writes_since_checkpoint,
7118 policy.checkpoint_write_interval + 2,
7119 "overlapping checkpoint finalizers erased later event debt"
7120 );
7121 assert!(retention.maintenance_pending);
7122 Ok(())
7123 })();
7124 assert!(
7125 result.is_ok(),
7126 "overlapping checkpoint test failed: {result:?}"
7127 );
7128 }
7129
7130 #[test]
7131 fn production_store_checkpoint_retries_and_rejects_stale_schema_or_binding() {
7132 let result = (|| -> Result<(), Box<dyn Error>> {
7133 let database = production_database()?;
7134 let runtime = instance(90)?;
7135 database.store.record_usage_for_instance(
7136 runtime,
7137 UsageInstanceOwner::McpProcess,
7138 &event("checkpoint", 100, 10),
7139 false,
7140 )?;
7141 let reader =
7142 AtlasStore::open_read_only_for_project(&database.database_path, &database.root)?;
7143 let _: i64 =
7144 reader
7145 .connection
7146 .query_row("SELECT COUNT(*) FROM usage_events", [], |row| row.get(0))?;
7147 database.store.record_usage_for_instance(
7148 runtime,
7149 UsageInstanceOwner::McpProcess,
7150 &event("checkpoint", 90, 10),
7151 false,
7152 )?;
7153 let policy = TelemetryRetentionPolicy::default();
7154 database.store.connection.execute(
7155 "UPDATE usage_retention_state
7156 SET writes_since_checkpoint = ?1 WHERE singleton = 1",
7157 [to_i64(
7158 "checkpoint_write_interval",
7159 policy.checkpoint_write_interval,
7160 )?],
7161 )?;
7162 maintain_after_commit_for_project(
7163 &database.store.connection,
7164 database.store.validated_project_root.as_deref(),
7165 database.project,
7166 policy,
7167 )?;
7168 let busy = database.store.telemetry_retention_state()?;
7169 assert_eq!(busy.checkpoint_state, TelemetryCheckpointState::Busy);
7170 assert_eq!(
7171 busy.writes_since_checkpoint,
7172 policy.checkpoint_write_interval
7173 );
7174
7175 reader.finish_index_read_snapshot()?;
7176 maintain_after_commit_for_project(
7177 &database.store.connection,
7178 database.store.validated_project_root.as_deref(),
7179 database.project,
7180 policy,
7181 )?;
7182 let completed = database.store.telemetry_retention_state()?;
7183 assert_eq!(
7184 completed.checkpoint_state,
7185 TelemetryCheckpointState::Completed
7186 );
7187 assert_eq!(completed.writes_since_checkpoint, 0);
7188
7189 database.store.connection.execute(
7190 "UPDATE usage_retention_state
7191 SET writes_since_checkpoint = ?1 WHERE singleton = 1",
7192 [to_i64(
7193 "checkpoint_write_interval",
7194 policy.checkpoint_write_interval,
7195 )?],
7196 )?;
7197 let newer_owner = Connection::open(&database.database_path)?;
7198 let future_schema = crate::schema::SCHEMA_VERSION + 1;
7199 newer_owner.execute(
7200 "UPDATE metadata SET value = ?2 WHERE key = ?1",
7201 params![crate::schema::SCHEMA_VERSION_KEY, future_schema.to_string()],
7202 )?;
7203 let wal_path = crate::schema::sqlite_sidecar_path(&database.database_path, "-wal");
7204 let wal_before = fs::read(&wal_path)?;
7205 let maintenance_before = database.store.telemetry_retention_state()?;
7206 let Err(error) = maintain_after_commit_for_project(
7207 &database.store.connection,
7208 database.store.validated_project_root.as_deref(),
7209 database.project,
7210 policy,
7211 ) else {
7212 return Err(std::io::Error::other(
7213 "post-commit maintenance accepted a newer schema",
7214 )
7215 .into());
7216 };
7217 assert!(matches!(
7218 error,
7219 DbError::SchemaVersion { found, expected }
7220 if found == future_schema && expected == crate::schema::SCHEMA_VERSION
7221 ));
7222 assert_eq!(
7223 database.store.telemetry_retention_state()?,
7224 maintenance_before,
7225 "newer-schema maintenance refusal changed retention state"
7226 );
7227 assert_eq!(
7228 fs::read(&wal_path)?,
7229 wal_before,
7230 "newer-schema maintenance refusal checkpointed the WAL"
7231 );
7232 assert_eq!(
7233 newer_owner.query_row(
7234 "SELECT value FROM metadata WHERE key = ?1",
7235 [crate::schema::SCHEMA_VERSION_KEY],
7236 |row| row.get::<_, String>(0),
7237 )?,
7238 future_schema.to_string(),
7239 "newer-schema owner stopped observing its schema"
7240 );
7241 newer_owner.execute(
7242 "UPDATE metadata SET value = ?2 WHERE key = ?1",
7243 params![
7244 crate::schema::SCHEMA_VERSION_KEY,
7245 crate::schema::SCHEMA_VERSION.to_string()
7246 ],
7247 )?;
7248
7249 database.store.connection.execute(
7250 "UPDATE usage_retention_state
7251 SET writes_since_checkpoint = ?1 WHERE singleton = 1",
7252 [to_i64(
7253 "checkpoint_write_interval",
7254 policy.checkpoint_write_interval,
7255 )?],
7256 )?;
7257 let transition_before = database.store.telemetry_retention_state()?;
7258 let Err(error) = maintain_after_commit_for_project_with_checkpoint(
7259 &database.store.connection,
7260 database.store.validated_project_root.as_deref(),
7261 database.project,
7262 policy,
7263 |connection| {
7264 let state = passive_checkpoint_state(connection);
7265 newer_owner.execute(
7266 "UPDATE metadata SET value = ?2 WHERE key = ?1",
7267 params![crate::schema::SCHEMA_VERSION_KEY, future_schema.to_string()],
7268 )?;
7269 Ok(state)
7270 },
7271 ) else {
7272 return Err(std::io::Error::other(
7273 "maintenance finalized after a schema transition at the checkpoint boundary",
7274 )
7275 .into());
7276 };
7277 assert!(matches!(
7278 error,
7279 DbError::SchemaVersion { found, expected }
7280 if found == future_schema && expected == crate::schema::SCHEMA_VERSION
7281 ));
7282 assert_eq!(
7283 database.store.telemetry_retention_state()?,
7284 transition_before,
7285 "checkpoint-boundary schema transition changed retention state"
7286 );
7287 assert_eq!(
7288 newer_owner.query_row(
7289 "SELECT value FROM metadata WHERE key = ?1",
7290 [crate::schema::SCHEMA_VERSION_KEY],
7291 |row| row.get::<_, String>(0),
7292 )?,
7293 future_schema.to_string(),
7294 "checkpoint-boundary refusal changed the concurrent owner's schema"
7295 );
7296 newer_owner.execute(
7297 "UPDATE metadata SET value = ?2 WHERE key = ?1",
7298 params![
7299 crate::schema::SCHEMA_VERSION_KEY,
7300 crate::schema::SCHEMA_VERSION.to_string()
7301 ],
7302 )?;
7303
7304 let old_project = database.project;
7305 let captured = database.store.captured_project_binding()?;
7306 assert_eq!(captured.project_instance_id, old_project);
7307 let detached = AtlasStore::transition_project_root(
7308 &database.database_path,
7309 &database.root,
7310 crate::ProjectRootTransition::Detach,
7311 )?;
7312 assert_ne!(detached.project_instance_id, old_project);
7313 assert!(
7314 database
7315 .store
7316 .revalidate_captured_project_binding()
7317 .is_err()
7318 );
7319 let stale = maintain_after_commit_for_project(
7320 &database.store.connection,
7321 database.store.validated_project_root.as_deref(),
7322 old_project,
7323 policy,
7324 );
7325 assert!(stale.is_err());
7326 Ok(())
7327 })();
7328 assert!(
7329 result.is_ok(),
7330 "production checkpoint lifecycle test failed: {result:?}"
7331 );
7332 }
7333
7334 #[test]
7335 fn production_store_reuses_freed_raw_pages_without_request_path_vacuum() {
7336 let result = (|| -> Result<(), Box<dyn Error>> {
7337 let database = production_database()?;
7338 let runtime = instance(91)?;
7339 let generous = TelemetryRetentionPolicy {
7340 max_raw_rows: 200,
7341 max_raw_logical_bytes: 2 * 1_024 * 1_024,
7342 checkpoint_write_interval: usize::MAX,
7343 ..TelemetryRetentionPolicy::default()
7344 };
7345 let mut large = event("page-reuse", 100, 10);
7346 large.query = Some("x".repeat(3_500));
7347 for now in 10_000..10_128 {
7348 record_transaction_at(
7349 &database.store.connection,
7350 database.project,
7351 runtime,
7352 UsageInstanceOwner::McpProcess,
7353 &large,
7354 generous,
7355 false,
7356 now,
7357 )?;
7358 }
7359 let allocated = database.store.telemetry_retention_state()?;
7360 let compact = TelemetryRetentionPolicy {
7361 max_raw_rows: 16,
7362 checkpoint_write_interval: usize::MAX,
7363 ..generous
7364 };
7365 let transaction = Transaction::new_unchecked(
7366 &database.store.connection,
7367 TransactionBehavior::Immediate,
7368 )?;
7369 let mut pruned = 0usize;
7370 loop {
7371 let deleted = prune_raw_once(&transaction, compact, 10_128)?;
7372 pruned = pruned
7373 .checked_add(deleted)
7374 .ok_or(DbError::TelemetryIntegerOverflow {
7375 field: "pruned_raw_rows",
7376 })?;
7377 if deleted == 0 {
7378 break;
7379 }
7380 }
7381 refresh_retention_state(&transaction, compact, 10_128, pruned, 0, 0, 0)?;
7382 transaction.commit()?;
7383 let after_prune = database.store.telemetry_retention_state()?;
7384 assert_eq!(after_prune.raw_rows, 16);
7385 assert!(after_prune.freelist_pages > 0);
7386 assert_eq!(after_prune.page_count, allocated.page_count);
7387
7388 let refill = TelemetryRetentionPolicy {
7389 max_raw_rows: 80,
7390 ..compact
7391 };
7392 for now in 20_000..20_064 {
7393 record_transaction_at(
7394 &database.store.connection,
7395 database.project,
7396 runtime,
7397 UsageInstanceOwner::McpProcess,
7398 &large,
7399 refill,
7400 false,
7401 now,
7402 )?;
7403 }
7404 let after_refill = database.store.telemetry_retention_state()?;
7405 assert_eq!(after_refill.raw_rows, 80);
7406 assert!(after_refill.page_count <= allocated.page_count);
7407 assert!(after_refill.freelist_pages < after_prune.freelist_pages);
7408 Ok(())
7409 })();
7410 assert!(
7411 result.is_ok(),
7412 "production page-reuse test failed: {result:?}"
7413 );
7414 }
7415
7416 #[test]
7417 fn synchronization_coalesces_dimensions_that_share_an_overflow_bucket()
7418 -> Result<(), Box<dyn Error>> {
7419 let temp = tempfile::tempdir()?;
7420 let control_root = temp.path().join("control");
7421 let worktree_root = temp.path().join("worktree");
7422 let common = temp.path().join("common.git");
7423 let administrative = common.join("worktrees/overflow");
7424 for path in [&control_root, &worktree_root, &administrative] {
7425 fs::create_dir_all(path)?;
7426 }
7427 let control = store_at(&control_root)?;
7428 let worktree = store_at(&worktree_root)?;
7429 let project = worktree
7430 .validated_project_instance_id
7431 .ok_or(DbError::ProjectInstanceIdentityMissing)?;
7432 let alias = WorktreeAlias::parse("overflow")?;
7433 let registration = control.register_worktree(
7434 &alias,
7435 &common,
7436 &administrative,
7437 &"66".repeat(32),
7438 &worktree_root,
7439 Some(project),
7440 10,
7441 )?;
7442 let policy = TelemetryRetentionPolicy::default().validate()?;
7443 for index in 0..policy.max_dimensions {
7444 if retention_counter(&control.connection, RetentionCounter::DimensionRows)?
7445 >= policy.max_dimensions
7446 {
7447 break;
7448 }
7449 let mut dimension = DimensionValues::from_event(&event("control", 1, 0));
7450 dimension.model = format!("control-{index}");
7451 ensure_dimension(&control.connection, &dimension, policy)?;
7452 }
7453 require_eq(
7454 &retention_counter(&control.connection, RetentionCounter::DimensionRows)?,
7455 &policy.max_dimensions,
7456 "full control dimension capacity",
7457 )?;
7458
7459 let first = event("local", 100, 20);
7460 let mut second = event("local", 80, 20);
7461 second.provider = "other-provider".to_string();
7462 second.baseline_identity = "source:src/other.rs".to_string();
7463 second.baseline_fingerprint = "source:src/other.rs:v1".to_string();
7464 for usage in [&first, &second] {
7465 worktree.record_usage(usage)?;
7466 }
7467 require_eq(
7468 &control
7469 .synchronize_worktree_usage(&alias, &worktree.export_worktree_usage_snapshot()?)?,
7470 &WorktreeUsageSyncState::Synchronized,
7471 "first overflow synchronization",
7472 )?;
7473 for usage in [&first, &second] {
7474 worktree.record_usage(usage)?;
7475 }
7476 require_eq(
7477 &control
7478 .synchronize_worktree_usage(&alias, &worktree.export_worktree_usage_snapshot()?)?,
7479 &WorktreeUsageSyncState::Synchronized,
7480 "monotonic overflow synchronization",
7481 )?;
7482 require_eq(
7483 &control.registered_worktree_token_overview(&alias)?.calls,
7484 &4,
7485 "coalesced overflow lifetime calls",
7486 )?;
7487 let rows = control.connection.query_row(
7488 "SELECT SUM(day_epoch = -1), SUM(day_epoch >= 0)
7489 FROM worktree_usage_aggregates
7490 WHERE registration_id = ?1 AND source_kind = ?2",
7491 params![registration.registration_id, WORKTREE_USAGE_SYNCHRONIZED],
7492 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
7493 )?;
7494 require_eq(&rows, &(1, 1), "coalesced overflow aggregate rows")?;
7495 Ok(())
7496 }
7497
7498 #[test]
7499 fn routed_and_synchronized_worktree_usage_remain_exact_monotonic_and_retained()
7500 -> Result<(), Box<dyn Error>> {
7501 let temp = tempfile::tempdir()?;
7502 let control_root = temp.path().join("control");
7503 let worktree_root = temp.path().join("feature");
7504 let other_root = temp.path().join("other");
7505 let common = temp.path().join("common.git");
7506 let administrative = common.join("worktrees/feature");
7507 for path in [&control_root, &worktree_root, &other_root, &administrative] {
7508 fs::create_dir_all(path)?;
7509 }
7510 let control = store_at(&control_root)?;
7511 let worktree = store_at(&worktree_root)?;
7512 let other = store_at(&other_root)?;
7513 let worktree_project = worktree
7514 .validated_project_instance_id
7515 .ok_or(DbError::ProjectInstanceIdentityMissing)?;
7516 let alias = WorktreeAlias::parse("issue-430")?;
7517 let registration = control.register_worktree(
7518 &alias,
7519 &common,
7520 &administrative,
7521 &"11".repeat(32),
7522 &worktree_root,
7523 Some(worktree_project),
7524 10,
7525 )?;
7526
7527 let routed_instance = instance(41)?;
7528 control.record_usage_for_worktree_instance(
7529 routed_instance,
7530 UsageInstanceOwner::McpProcess,
7531 registration.registration_id,
7532 &event("routed", 100, 20),
7533 false,
7534 )?;
7535 require_eq(
7536 &control.token_overview(None)?.calls,
7537 &1,
7538 "routed native total",
7539 )?;
7540 require_eq(
7541 &control.repository_token_overview()?.calls,
7542 &1,
7543 "routed repository total",
7544 )?;
7545 require_eq(
7546 &control.registered_worktree_token_overview(&alias)?.calls,
7547 &1,
7548 "routed origin total",
7549 )?;
7550 require(
7551 matches!(
7552 control.record_usage_for_instance(
7553 routed_instance,
7554 UsageInstanceOwner::McpProcess,
7555 &event("routed", 100, 20),
7556 false,
7557 ),
7558 Err(DbError::WorktreeTelemetryOriginConflict)
7559 ),
7560 "runtime instance crossed telemetry origins",
7561 )?;
7562 require_eq(
7563 &control.repository_token_overview()?.calls,
7564 &1,
7565 "origin-conflict rollback total",
7566 )?;
7567
7568 worktree.record_usage(&event("local", 80, 20))?;
7569 let first = worktree.export_worktree_usage_snapshot()?;
7570 require_eq(&first.revision(), &1, "first local revision")?;
7571 require(
7572 first.row_count() >= 2,
7573 "first snapshot omitted aggregate rows",
7574 )?;
7575 require(
7576 first.logical_bytes() > 0,
7577 "first snapshot omitted logical bytes",
7578 )?;
7579 require_eq(
7580 &control.synchronize_worktree_usage(&alias, &first)?,
7581 &WorktreeUsageSyncState::Synchronized,
7582 "first synchronization",
7583 )?;
7584 require_eq(
7585 &control.repository_token_overview()?.calls,
7586 &2,
7587 "first combined total",
7588 )?;
7589 require_eq(
7590 &control.registered_worktree_token_overview(&alias)?.calls,
7591 &2,
7592 "first exact worktree total",
7593 )?;
7594 require_eq(
7595 &control.synchronize_worktree_usage(&alias, &first)?,
7596 &WorktreeUsageSyncState::Current,
7597 "stale synchronization",
7598 )?;
7599 require_eq(
7600 &control.repository_token_overview()?.calls,
7601 &2,
7602 "stale synchronization total",
7603 )?;
7604
7605 worktree.record_usage(&event("local", 90, 20))?;
7606 let second = worktree.export_worktree_usage_snapshot()?;
7607 require_eq(&second.revision(), &2, "second local revision")?;
7608 require_eq(
7609 &control.synchronize_worktree_usage(&alias, &second)?,
7610 &WorktreeUsageSyncState::Synchronized,
7611 "second synchronization",
7612 )?;
7613 require_eq(
7614 &control.repository_token_overview()?.calls,
7615 &3,
7616 "second combined total",
7617 )?;
7618 require_eq(
7619 &control.registered_worktree_token_overview(&alias)?.calls,
7620 &3,
7621 "second exact worktree total",
7622 )?;
7623
7624 let first_lifetime = first
7625 .rows
7626 .iter()
7627 .find(|row| row.day_epoch == -1)
7628 .cloned()
7629 .ok_or_else(|| io::Error::other("first lifetime row missing"))?;
7630 let first_daily = first
7631 .rows
7632 .iter()
7633 .find(|row| row.day_epoch >= 0 && row.dimension_id == first_lifetime.dimension_id)
7634 .cloned()
7635 .ok_or_else(|| io::Error::other("first daily row missing"))?;
7636 let mut daily_rollback = second.clone();
7637 daily_rollback.revision = 3;
7638 let newer_lifetime = daily_rollback
7639 .rows
7640 .iter_mut()
7641 .find(|row| row.day_epoch == -1 && row.dimension_id == first_lifetime.dimension_id)
7642 .ok_or_else(|| io::Error::other("newer lifetime row missing"))?;
7643 newer_lifetime.counters = newer_lifetime
7644 .counters
7645 .checked_add(first_lifetime.counters)?;
7646 let rolled_back_daily = daily_rollback
7647 .rows
7648 .iter_mut()
7649 .find(|row| {
7650 row.day_epoch == first_daily.day_epoch
7651 && row.dimension_id == first_daily.dimension_id
7652 })
7653 .ok_or_else(|| io::Error::other("newer daily row missing"))?;
7654 rolled_back_daily.counters = first_daily.counters;
7655 daily_rollback.logical_bytes = validate_worktree_usage_snapshot(
7656 &daily_rollback.dimensions,
7657 &daily_rollback.rows,
7658 TelemetryRetentionPolicy::default().validate()?,
7659 )?;
7660 require(
7661 matches!(
7662 control.synchronize_worktree_usage(&alias, &daily_rollback),
7663 Err(DbError::WorktreeRegistrationRow {
7664 reason: "aggregate snapshot reduces accepted retained daily totals"
7665 })
7666 ),
7667 "newer revision from a recent backup reduced an accepted daily bucket",
7668 )?;
7669 require_eq(
7670 &control
7671 .worktree_registration(&alias)?
7672 .accepted_telemetry_revision,
7673 &2,
7674 "revision after restored daily snapshot",
7675 )?;
7676 let retained_daily_calls = control.connection.query_row(
7677 "SELECT calls FROM worktree_usage_aggregates
7678 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch >= 0
7679 ORDER BY day_epoch, dimension_id LIMIT 1",
7680 params![registration.registration_id, WORKTREE_USAGE_SYNCHRONIZED],
7681 |row| row.get::<_, i64>(0),
7682 )?;
7683 require_eq(
7684 &retained_daily_calls,
7685 &2,
7686 "daily bucket after restored daily snapshot",
7687 )?;
7688
7689 let mut restored = first;
7690 restored.revision = 3;
7691 let source_dimension = restored
7692 .dimensions
7693 .values()
7694 .next()
7695 .cloned()
7696 .ok_or_else(|| io::Error::other("restored source dimension missing"))?;
7697 let compensating_counters = second
7698 .rows
7699 .iter()
7700 .filter(|row| row.day_epoch == -1)
7701 .try_fold(AggregateCounters::default(), |total, row| {
7702 total.checked_add(row.counters)
7703 })?;
7704 let compensating_dimension_id = restored
7705 .dimensions
7706 .keys()
7707 .next_back()
7708 .copied()
7709 .unwrap_or_default()
7710 .checked_add(1)
7711 .ok_or_else(|| io::Error::other("restored dimension identifier overflow"))?;
7712 let mut compensating_dimension = source_dimension;
7713 compensating_dimension.model.push_str("-other");
7714 restored
7715 .dimensions
7716 .insert(compensating_dimension_id, compensating_dimension);
7717 restored.rows.push(WorktreeUsageSnapshotRow {
7718 day_epoch: -1,
7719 dimension_id: compensating_dimension_id,
7720 counters: compensating_counters,
7721 });
7722 restored.logical_bytes = validate_worktree_usage_snapshot(
7723 &restored.dimensions,
7724 &restored.rows,
7725 TelemetryRetentionPolicy::default().validate()?,
7726 )?;
7727 require(
7728 matches!(
7729 control.synchronize_worktree_usage(&alias, &restored),
7730 Err(DbError::WorktreeRegistrationRow {
7731 reason: "aggregate snapshot reduces accepted lifetime totals"
7732 })
7733 ),
7734 "newer revision from an older backup reduced accepted lifetime totals",
7735 )?;
7736 require_eq(
7737 &control
7738 .worktree_registration(&alias)?
7739 .accepted_telemetry_revision,
7740 &2,
7741 "revision after restored older snapshot",
7742 )?;
7743 require_eq(
7744 &control.repository_token_overview()?.calls,
7745 &3,
7746 "total after restored older snapshot",
7747 )?;
7748
7749 let mut invalid = second.clone();
7750 invalid.revision = 3;
7751 invalid
7752 .rows
7753 .first_mut()
7754 .ok_or_else(|| std::io::Error::other("invalid snapshot row missing"))?
7755 .counters
7756 .calls = -1;
7757 require(
7758 matches!(
7759 control.synchronize_worktree_usage(&alias, &invalid),
7760 Err(DbError::TelemetryIntegerOverflow { field: "calls" })
7761 ),
7762 "invalid snapshot was not rejected",
7763 )?;
7764 require_eq(
7765 &control
7766 .worktree_registration(&alias)?
7767 .accepted_telemetry_revision,
7768 &2,
7769 "revision after invalid snapshot",
7770 )?;
7771 require_eq(
7772 &control.repository_token_overview()?.calls,
7773 &3,
7774 "total after invalid snapshot",
7775 )?;
7776
7777 let other_snapshot = other.export_worktree_usage_snapshot()?;
7778 require(
7779 matches!(
7780 control.synchronize_worktree_usage(&alias, &other_snapshot),
7781 Err(DbError::WorktreeTelemetryProjectMismatch { .. })
7782 ),
7783 "mismatched project snapshot was not rejected",
7784 )?;
7785 let aged_rows = control.connection.execute(
7786 "UPDATE worktree_usage_aggregates SET day_epoch = 0
7787 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch >= 0",
7788 params![registration.registration_id, WORKTREE_USAGE_SYNCHRONIZED],
7789 )?;
7790 require_eq(&aged_rows, &1, "daily rows prepared for expiry")?;
7791 let mut expired = second;
7792 expired.revision = 3;
7793 expired.rows.retain(|row| row.day_epoch == -1);
7794 expired.logical_bytes = validate_worktree_usage_snapshot(
7795 &expired.dimensions,
7796 &expired.rows,
7797 TelemetryRetentionPolicy::default().validate()?,
7798 )?;
7799 require_eq(
7800 &control.synchronize_worktree_usage(&alias, &expired)?,
7801 &WorktreeUsageSyncState::Synchronized,
7802 "expired daily bucket synchronization",
7803 )?;
7804 let expired_rows = control.connection.query_row(
7805 "SELECT COUNT(*) FROM worktree_usage_aggregates
7806 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch = 0",
7807 params![registration.registration_id, WORKTREE_USAGE_SYNCHRONIZED],
7808 |row| row.get::<_, i64>(0),
7809 )?;
7810 require_eq(&expired_rows, &0, "expired synchronized daily rows")?;
7811 let contending_local = store_at(&worktree_root)?;
7812 let local_reader = AtlasStore::open_read_only_for_project(
7813 &worktree_root.join(".projectatlas").join("projectatlas.db"),
7814 &worktree_root,
7815 )?;
7816 let (retired, final_sync) =
7817 local_reader.with_exclusive_worktree_usage_snapshot(|snapshot| {
7818 let blocked = contending_local.record_usage(&event("after-export", 70, 20));
7819 if !blocked.as_ref().is_err_and(DbError::is_write_unavailable) {
7820 return Err(DbError::WorktreeRegistrationRow {
7821 reason: "local writer was not excluded during final synchronization",
7822 });
7823 }
7824 control.retire_worktree_with_usage_snapshot(
7825 registration.registration_id,
7826 &alias,
7827 &worktree_root,
7828 snapshot.project_instance_id(),
7829 snapshot,
7830 20,
7831 )
7832 })?;
7833 require_eq(
7834 &retired.state,
7835 &WorktreeRegistrationState::Retired,
7836 "writer-excluded retirement state",
7837 )?;
7838 require_eq(
7839 &final_sync,
7840 &WorktreeUsageSyncState::Current,
7841 "writer-excluded final synchronization",
7842 )?;
7843 contending_local.record_usage(&event("after-retirement", 70, 20))?;
7844 require_eq(
7845 &contending_local.token_overview(None)?.calls,
7846 &3,
7847 "local usage after retirement",
7848 )?;
7849 require_eq(
7850 &control.repository_token_overview()?.calls,
7851 &3,
7852 "retired repository total",
7853 )?;
7854 Ok(())
7855 }
7856
7857 #[test]
7858 fn worktree_usage_export_keeps_revision_and_aggregates_in_one_read_snapshot()
7859 -> Result<(), Box<dyn Error>> {
7860 let temp = tempfile::tempdir()?;
7861 let root = temp.path().join("worktree");
7862 fs::create_dir(&root)?;
7863 let database_path = root.join("projectatlas.db");
7864 let worktree = AtlasStore::open_for_project(&database_path, &root)?;
7865 worktree.record_usage(&event("snapshot", 100, 20))?;
7866
7867 let (entered_sender, entered_receiver) = sync_channel(1);
7868 let (resume_sender, resume_receiver) = sync_channel(1);
7869 WORKTREE_SNAPSHOT_EXPORT_BLOCKER.with(|slot| {
7870 *slot.borrow_mut() = Some(SnapshotExportBlocker {
7871 entered: entered_sender,
7872 resume: resume_receiver,
7873 });
7874 });
7875 worktree.connection.trace_v2(
7876 rusqlite::trace::TraceEventCodes::SQLITE_TRACE_STMT,
7877 Some(block_worktree_snapshot_aggregate_event),
7878 );
7879
7880 let writer_root = root.clone();
7881 let writer_database = database_path.clone();
7882 let writer = std::thread::spawn(move || {
7883 let result = (|| -> Result<(), String> {
7884 entered_receiver
7885 .recv_timeout(Duration::from_secs(10))
7886 .map_err(|error| error.to_string())?;
7887 let writer = AtlasStore::open_for_project(&writer_database, &writer_root)
7888 .map_err(|error| error.to_string())?;
7889 writer
7890 .record_usage(&event("snapshot", 90, 20))
7891 .map_err(|error| error.to_string())
7892 })();
7893 let _resume = resume_sender.send(());
7894 result
7895 });
7896
7897 let snapshot_result = worktree.export_worktree_usage_snapshot();
7898 worktree
7899 .connection
7900 .trace_v2(rusqlite::trace::TraceEventCodes::empty(), None);
7901 WORKTREE_SNAPSHOT_EXPORT_BLOCKER.with(|slot| {
7902 slot.borrow_mut().take();
7903 });
7904 writer
7905 .join()
7906 .map_err(|_panic| io::Error::other("snapshot writer panicked"))?
7907 .map_err(io::Error::other)?;
7908 let snapshot = snapshot_result?;
7909 require_eq(&snapshot.revision(), &1, "exported snapshot revision")?;
7910 let exported_calls = snapshot
7911 .rows
7912 .iter()
7913 .filter(|row| row.day_epoch == -1)
7914 .map(|row| row.counters.calls)
7915 .sum::<i64>();
7916 require_eq(&exported_calls, &1, "exported aggregate calls")?;
7917
7918 let current = worktree.export_worktree_usage_snapshot()?;
7919 require_eq(¤t.revision(), &2, "current snapshot revision")?;
7920 let current_calls = current
7921 .rows
7922 .iter()
7923 .filter(|row| row.day_epoch == -1)
7924 .map(|row| row.counters.calls)
7925 .sum::<i64>();
7926 require_eq(¤t_calls, &2, "current aggregate calls")?;
7927 let reader = AtlasStore::open_read_only_for_project(&database_path, &root)?;
7928 let read_only = reader.export_worktree_usage_snapshot()?;
7929 require_eq(&read_only.revision(), &2, "read-only snapshot revision")?;
7930 reader.finish_index_read_snapshot()?;
7931 Ok(())
7932 }
7933
7934 #[test]
7935 fn routed_daily_capacity_preserves_synchronized_snapshot() -> Result<(), Box<dyn Error>> {
7936 let temp = tempfile::tempdir()?;
7937 let control_root = temp.path().join("control");
7938 let worktree_root = temp.path().join("feature");
7939 let administrative = temp.path().join("common.git/worktrees/feature");
7940 for path in [&control_root, &worktree_root, &administrative] {
7941 fs::create_dir_all(path)?;
7942 }
7943 let control = store_at(&control_root)?;
7944 let worktree = store_at(&worktree_root)?;
7945 let worktree_project = worktree
7946 .validated_project_instance_id
7947 .ok_or(DbError::ProjectInstanceIdentityMissing)?;
7948 let alias = WorktreeAlias::parse("capacity")?;
7949 let registration = control.register_worktree(
7950 &alias,
7951 &temp.path().join("common.git"),
7952 &administrative,
7953 &"22".repeat(32),
7954 &worktree_root,
7955 Some(worktree_project),
7956 10,
7957 )?;
7958 let now = now_epoch_seconds()?;
7959 for (offset, instance_id) in [(2, 40), (1, 41)] {
7960 record_transaction_at(
7961 &worktree.connection,
7962 worktree_project,
7963 instance(instance_id)?,
7964 UsageInstanceOwner::CliInvocation,
7965 &event("synchronized", 100, 20),
7966 TelemetryRetentionPolicy::default(),
7967 true,
7968 now - i64::from(offset) * SECONDS_PER_DAY,
7969 )?;
7970 }
7971 let snapshot = worktree.export_worktree_usage_snapshot()?;
7972 require_eq(
7973 &control.synchronize_worktree_usage(&alias, &snapshot)?,
7974 &WorktreeUsageSyncState::Synchronized,
7975 "initial synchronized snapshot",
7976 )?;
7977
7978 let policy = TelemetryRetentionPolicy {
7979 max_daily_rows: 2,
7980 ..TelemetryRetentionPolicy::default()
7981 };
7982 let control_project = control
7983 .validated_project_instance_id
7984 .ok_or(DbError::ProjectInstanceIdentityMissing)?;
7985 let routed_instance = instance(42)?;
7986 let routed = control.with_validated_write(|connection| {
7987 record_usage_for_project(
7988 connection,
7989 control_project,
7990 routed_instance,
7991 UsageInstanceOwner::McpProcess,
7992 Some(registration.registration_id),
7993 &event("routed", 90, 20),
7994 policy,
7995 true,
7996 )
7997 });
7998 require(
7999 matches!(
8000 routed,
8001 Err(DbError::WorktreeTelemetrySnapshotLimit {
8002 resource: "daily_rows",
8003 limit: 2,
8004 observed: 3
8005 })
8006 ),
8007 "routed usage displaced a synchronized daily snapshot",
8008 )?;
8009 require_eq(
8010 &control.synchronize_worktree_usage(&alias, &snapshot)?,
8011 &WorktreeUsageSyncState::Current,
8012 "accepted snapshot revision after routed capacity failure",
8013 )?;
8014 require_eq(
8015 &control.registered_worktree_token_overview(&alias)?.calls,
8016 &2,
8017 "synchronized total after routed capacity failure",
8018 )?;
8019 let synchronized_daily = control.connection.query_row(
8020 "SELECT COUNT(*) FROM worktree_usage_aggregates
8021 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch >= 0",
8022 params![registration.registration_id, WORKTREE_USAGE_SYNCHRONIZED],
8023 |row| row.get::<_, i64>(0),
8024 )?;
8025 require_eq(&synchronized_daily, &2, "retained synchronized daily rows")?;
8026 Ok(())
8027 }
8028
8029 #[test]
8030 fn synchronization_prunes_expired_daily_rows_from_retired_origins() -> Result<(), Box<dyn Error>>
8031 {
8032 let temp = tempfile::tempdir()?;
8033 let control_root = temp.path().join("control");
8034 let retired_root = temp.path().join("retired");
8035 let current_root = temp.path().join("current");
8036 let common = temp.path().join("common.git");
8037 let retired_administrative = common.join("worktrees/retired");
8038 let current_administrative = common.join("worktrees/current");
8039 for path in [
8040 &control_root,
8041 &retired_root,
8042 ¤t_root,
8043 &retired_administrative,
8044 ¤t_administrative,
8045 ] {
8046 fs::create_dir_all(path)?;
8047 }
8048 let control = store_at(&control_root)?;
8049 let retired = store_at(&retired_root)?;
8050 let retired_project = retired
8051 .validated_project_instance_id
8052 .ok_or(DbError::ProjectInstanceIdentityMissing)?;
8053 let retired_alias = WorktreeAlias::parse("retired")?;
8054 let retired_registration = control.register_worktree(
8055 &retired_alias,
8056 &common,
8057 &retired_administrative,
8058 &"33".repeat(32),
8059 &retired_root,
8060 Some(retired_project),
8061 10,
8062 )?;
8063 retired.record_usage(&event("retired", 100, 20))?;
8064 let accepted = retired.export_worktree_usage_snapshot()?;
8065 require_eq(
8066 &control.synchronize_worktree_usage(&retired_alias, &accepted)?,
8067 &WorktreeUsageSyncState::Synchronized,
8068 "initial retired-origin synchronization",
8069 )?;
8070
8071 control.connection.execute(
8072 "DELETE FROM worktree_usage_aggregates
8073 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch >= 0",
8074 params![
8075 retired_registration.registration_id,
8076 WORKTREE_USAGE_SYNCHRONIZED
8077 ],
8078 )?;
8079 let policy = TelemetryRetentionPolicy::default().validate()?;
8080 let max_daily_rows = to_i64("max_daily_rows", policy.max_daily_rows)?;
8081 let seeded = control.connection.execute(
8082 "WITH RECURSIVE expired_days(day_epoch) AS (
8083 VALUES(0)
8084 UNION ALL
8085 SELECT day_epoch + 1 FROM expired_days WHERE day_epoch + 1 < ?3
8086 )
8087 INSERT INTO worktree_usage_aggregates(
8088 registration_id, source_kind, day_epoch, dimension_id,
8089 calls, estimated_without, estimated_with, observed_without,
8090 observed_with, modeled_without, modeled_with,
8091 deduped_modeled_without, deduped_modeled_with, repeated_baselines,
8092 observed_file_read_replacements, modeled_file_reads_avoided
8093 )
8094 SELECT aggregate.registration_id, aggregate.source_kind,
8095 expired_days.day_epoch, aggregate.dimension_id,
8096 aggregate.calls, aggregate.estimated_without, aggregate.estimated_with,
8097 aggregate.observed_without, aggregate.observed_with,
8098 aggregate.modeled_without, aggregate.modeled_with,
8099 aggregate.deduped_modeled_without, aggregate.deduped_modeled_with,
8100 aggregate.repeated_baselines, aggregate.observed_file_read_replacements,
8101 aggregate.modeled_file_reads_avoided
8102 FROM worktree_usage_aggregates AS aggregate
8103 CROSS JOIN expired_days
8104 WHERE aggregate.registration_id = ?1 AND aggregate.source_kind = ?2
8105 AND aggregate.day_epoch = -1",
8106 params![
8107 retired_registration.registration_id,
8108 WORKTREE_USAGE_SYNCHRONIZED,
8109 max_daily_rows
8110 ],
8111 )?;
8112 require_eq(&seeded, &policy.max_daily_rows, "seeded expired daily rows")?;
8113
8114 let mut reduced = accepted;
8115 reduced.revision = reduced.revision.saturating_add(1);
8116 let lifetime = reduced
8117 .rows
8118 .iter_mut()
8119 .find(|row| row.day_epoch == -1)
8120 .ok_or_else(|| io::Error::other("retired lifetime row missing"))?;
8121 lifetime.counters.calls = 0;
8122 reduced.logical_bytes =
8123 validate_worktree_usage_snapshot(&reduced.dimensions, &reduced.rows, policy)?;
8124 require(
8125 matches!(
8126 control.synchronize_worktree_usage(&retired_alias, &reduced),
8127 Err(DbError::WorktreeRegistrationRow {
8128 reason: "aggregate snapshot reduces accepted lifetime totals"
8129 })
8130 ),
8131 "failed synchronization committed expired-row pruning",
8132 )?;
8133 let after_rollback = control.connection.query_row(
8134 "SELECT COUNT(*) FROM worktree_usage_aggregates
8135 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch >= 0",
8136 params![
8137 retired_registration.registration_id,
8138 WORKTREE_USAGE_SYNCHRONIZED
8139 ],
8140 |row| row.get::<_, i64>(0),
8141 )?;
8142 require_eq(
8143 &after_rollback,
8144 &max_daily_rows,
8145 "expired rows after failed synchronization",
8146 )?;
8147 control.retire_worktree(retired_registration.registration_id, &retired_alias, 20)?;
8148
8149 let current = store_at(¤t_root)?;
8150 let current_project = current
8151 .validated_project_instance_id
8152 .ok_or(DbError::ProjectInstanceIdentityMissing)?;
8153 let current_alias = WorktreeAlias::parse("current")?;
8154 let current_registration = control.register_worktree(
8155 ¤t_alias,
8156 &common,
8157 ¤t_administrative,
8158 &"44".repeat(32),
8159 ¤t_root,
8160 Some(current_project),
8161 30,
8162 )?;
8163 current.record_usage(&event("current", 90, 20))?;
8164 let mut current_snapshot = current.export_worktree_usage_snapshot()?;
8165 let mut expired = current_snapshot
8166 .rows
8167 .iter()
8168 .find(|row| row.day_epoch >= 0)
8169 .cloned()
8170 .ok_or_else(|| io::Error::other("current daily row missing"))?;
8171 expired.day_epoch = 0;
8172 current_snapshot
8173 .rows
8174 .extend(std::iter::repeat_n(expired, policy.max_daily_rows));
8175 current_snapshot.logical_bytes = validate_worktree_usage_snapshot(
8176 ¤t_snapshot.dimensions,
8177 ¤t_snapshot.rows,
8178 policy,
8179 )?;
8180 require_eq(
8181 &control.synchronize_worktree_usage(¤t_alias, ¤t_snapshot)?,
8182 &WorktreeUsageSyncState::Synchronized,
8183 "current-origin synchronization after global pruning",
8184 )?;
8185
8186 let retired_daily = control.connection.query_row(
8187 "SELECT COUNT(*) FROM worktree_usage_aggregates
8188 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch >= 0",
8189 params![
8190 retired_registration.registration_id,
8191 WORKTREE_USAGE_SYNCHRONIZED
8192 ],
8193 |row| row.get::<_, i64>(0),
8194 )?;
8195 require_eq(&retired_daily, &0, "retired expired daily rows")?;
8196 let lifetime_calls = control.connection.query_row(
8197 "SELECT SUM(calls) FROM worktree_usage_aggregates
8198 WHERE source_kind = ?1 AND day_epoch = -1",
8199 [WORKTREE_USAGE_SYNCHRONIZED],
8200 |row| row.get::<_, i64>(0),
8201 )?;
8202 require_eq(&lifetime_calls, &2, "synchronized lifetime truth")?;
8203 let current_daily_calls = control.connection.query_row(
8204 "SELECT SUM(calls) FROM worktree_usage_aggregates
8205 WHERE registration_id = ?1 AND source_kind = ?2 AND day_epoch >= 0",
8206 params![
8207 current_registration.registration_id,
8208 WORKTREE_USAGE_SYNCHRONIZED
8209 ],
8210 |row| row.get::<_, i64>(0),
8211 )?;
8212 require_eq(¤t_daily_calls, &1, "retained current daily truth")?;
8213 require_eq(
8214 &control.repository_token_overview()?.calls,
8215 &2,
8216 "repository lifetime truth",
8217 )?;
8218 Ok(())
8219 }
8220
8221 #[test]
8222 fn worktree_continuity_high_registration_aggregate_has_bounded_sql_and_rows()
8223 -> Result<(), Box<dyn Error>> {
8224 const ORIGINS: usize = 128;
8225 const EXPECTED_STATEMENTS: usize = 11_407;
8226 const EXPECTED_CHANGED_ROWS: u64 = 3_203;
8227
8228 let temp = tempfile::tempdir()?;
8229 let control_root = temp.path().join("control");
8230 fs::create_dir_all(&control_root)?;
8231 let database = control_root.join("projectatlas.db");
8232 let control = AtlasStore::open_for_project(&database, &control_root)?;
8233 let project = control.captured_project_binding()?.project_instance_id;
8234 let common = temp.path().join("common.git");
8235 let event = event("worktree-scale", 100, 20);
8236
8237 WORKTREE_TRACE.with(|statements| statements.borrow_mut().clear());
8238 control.connection.trace_v2(
8239 rusqlite::trace::TraceEventCodes::SQLITE_TRACE_STMT,
8240 Some(record_worktree_event),
8241 );
8242 let changed_before = control.connection.total_changes();
8243 for index in 0..ORIGINS {
8244 let suffix = index + 1;
8245 let alias = WorktreeAlias::parse(&format!("worktree-{suffix:03}"))?;
8246 let registration = control.register_worktree(
8247 &alias,
8248 &common,
8249 &common.join(format!("worktrees/{suffix:03}")),
8250 &format!("{suffix:064x}"),
8251 &temp.path().join(format!("worktree-{suffix:03}")),
8252 Some(ProjectInstanceId::from_bytes([u8::try_from(suffix)?; 16])?),
8253 u64::try_from(suffix)?,
8254 )?;
8255 record_usage_for_project(
8256 &control.connection,
8257 project,
8258 instance(u8::try_from(suffix)?)?,
8259 UsageInstanceOwner::McpProcess,
8260 Some(registration.registration_id),
8261 &event,
8262 TelemetryRetentionPolicy::default(),
8263 true,
8264 )?;
8265 }
8266 let registrations = control.worktree_registrations(false)?;
8267 let overview = control.repository_token_overview()?;
8268 control
8269 .connection
8270 .trace_v2(rusqlite::trace::TraceEventCodes::empty(), None);
8271 let changed_rows = control.connection.total_changes() - changed_before;
8272 let statements =
8273 WORKTREE_TRACE.with(|statements| std::mem::take(&mut *statements.borrow_mut()));
8274
8275 require_eq(®istrations.len(), &ORIGINS, "active registration count")?;
8276 require_eq(&overview.calls, &ORIGINS, "repository aggregate calls")?;
8277 require_eq(
8278 &statements.len(),
8279 &EXPECTED_STATEMENTS,
8280 "worktree scale statement count",
8281 )?;
8282 require_eq(
8283 &changed_rows,
8284 &EXPECTED_CHANGED_ROWS,
8285 "worktree scale changed rows",
8286 )?;
8287 Ok(())
8288 }
8289
8290 #[test]
8291 fn hot_telemetry_queries_use_owned_indexes_without_duplicate_primary_key_indexes() {
8292 let result = (|| -> Result<(), Box<dyn Error>> {
8293 let database = test_database()?;
8294 assert_plan_uses(
8295 &query_plan(
8296 &database.connection,
8297 "SELECT instance_row_id FROM usage_instances
8298 WHERE project_instance_id = X'01010101010101010101010101010101'
8299 AND runtime_instance_id = X'02020202020202020202020202020202'",
8300 )?,
8301 "sqlite_autoindex_usage_instances_1",
8302 );
8303 assert_plan_uses(
8304 &query_plan(
8305 &database.connection,
8306 "SELECT maximum_without FROM usage_instance_baselines
8307 WHERE instance_row_id = 1
8308 AND baseline_key = zeroblob(32)",
8309 )?,
8310 "PRIMARY KEY",
8311 );
8312 assert_plan_uses(
8313 &query_plan(
8314 &database.connection,
8315 "SELECT id FROM usage_events
8316 WHERE created_at_epoch < 1
8317 ORDER BY created_at_epoch, id LIMIT 1",
8318 )?,
8319 "idx_usage_created_at",
8320 );
8321 assert_plan_uses(
8322 &query_plan(
8323 &database.connection,
8324 "SELECT calls FROM usage_global_aggregates
8325 WHERE project_instance_id = X'01010101010101010101010101010101'
8326 AND dimension_id = 1",
8327 )?,
8328 "PRIMARY KEY",
8329 );
8330 assert_plan_uses(
8331 &query_plan(
8332 &database.connection,
8333 "SELECT day_epoch, dimension_id FROM worktree_usage_aggregates
8334 WHERE registration_id = 1 AND source_kind = 'synchronized'
8335 AND day_epoch >= 0
8336 ORDER BY day_epoch, dimension_id",
8337 )?,
8338 "PRIMARY KEY",
8339 );
8340 assert_plan_uses(
8341 &query_plan(
8342 &database.connection,
8343 "SELECT instance_row_id FROM usage_instances
8344 WHERE project_instance_id = X'01010101010101010101010101010101'
8345 AND caller_label = 'agent' AND state = 'active'
8346 ORDER BY started_at_epoch, instance_row_id LIMIT 1",
8347 )?,
8348 "idx_usage_instances_label_state",
8349 );
8350 assert_plan_uses(
8351 &query_plan(
8352 &database.connection,
8353 "SELECT project_instance_id FROM usage_daily_aggregates
8354 WHERE day_epoch < 1 ORDER BY day_epoch LIMIT 1",
8355 )?,
8356 "idx_usage_daily_retention",
8357 );
8358 assert_plan_uses(
8359 &query_plan(
8360 &database.connection,
8361 "SELECT project_instance_id FROM usage_label_tombstones
8362 WHERE expired_at_epoch < 1 ORDER BY expired_at_epoch LIMIT 1",
8363 )?,
8364 "idx_usage_label_tombstones_retention",
8365 );
8366 for redundant in [
8367 "idx_usage_baselines_active",
8368 "idx_usage_daily_range",
8369 "idx_usage_instance_daily_range",
8370 ] {
8371 let exists = database.connection.query_row(
8372 "SELECT EXISTS(SELECT 1 FROM sqlite_schema WHERE type = 'index' AND name = ?1)",
8373 [redundant],
8374 |row| row.get::<_, i64>(0),
8375 )?;
8376 assert_eq!(exists, 0, "redundant index {redundant} must stay absent");
8377 }
8378 Ok(())
8379 })();
8380 assert!(
8381 result.is_ok(),
8382 "telemetry query-plan test failed: {result:?}"
8383 );
8384 }
8385}