feat(usage): add local usage storage [1 of 4]

codex-rs/protocol/src/protocol.rs

@@ -1901,6 +1901,47 @@ pub struct TokenUsage {
     pub total_tokens: i64,
 }
 
+#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, Hash, JsonSchema, TS)]
+#[serde(rename_all = "snake_case")]
+#[ts(rename_all = "snake_case")]
+pub enum UsageContributorKind {
+    Skill,
+    Subagent,
+    AgentTask,
+    App,
+    McpServer,
+    Plugin,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]
+pub struct UsageContributor {
+    pub kind: UsageContributorKind,
+    pub id: String,
+    pub label: String,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]
+pub struct UsageAttributionContributor {
+    pub contributor: UsageContributor,
+    #[ts(type = "number")]
+    pub source_estimated_tokens: i64,
+    #[ts(type = "number")]
+    pub attributed_tokens: i64,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]
+pub struct UsageAttributionItem {
+    pub sample_id: String,
+    pub turn_id: String,
+    pub response_id: String,
+    #[ts(type = "number")]
+    pub occurred_at: i64,
+    pub token_usage: TokenUsage,
+    #[ts(type = "number")]
+    pub prompt_estimated_tokens: i64,
+    pub contributors: Vec<UsageAttributionContributor>,
+}
+
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]
 pub struct TokenUsageInfo {
     pub total_token_usage: TokenUsage,

codex-rs/state/migrations/0035_usage_samples.sql

@@ -0,0 +1,29 @@
+CREATE TABLE usage_samples (
+    sample_id TEXT PRIMARY KEY,
+    thread_id TEXT NOT NULL REFERENCES threads(id) ON DELETE CASCADE,
+    turn_id TEXT NOT NULL,
+    response_id TEXT NOT NULL,
+    occurred_at INTEGER NOT NULL,
+    input_tokens INTEGER NOT NULL,
+    cached_input_tokens INTEGER NOT NULL,
+    non_cached_input_tokens INTEGER NOT NULL,
+    output_tokens INTEGER NOT NULL,
+    reasoning_output_tokens INTEGER NOT NULL,
+    total_tokens INTEGER NOT NULL,
+    blended_tokens INTEGER NOT NULL,
+    prompt_estimated_tokens INTEGER NOT NULL
+);
+
+CREATE TABLE usage_sample_contributors (
+    sample_id TEXT NOT NULL REFERENCES usage_samples(sample_id) ON DELETE CASCADE,
+    kind TEXT NOT NULL,
+    contributor_id TEXT NOT NULL,
+    label TEXT NOT NULL,
+    source_estimated_tokens INTEGER NOT NULL,
+    attributed_tokens INTEGER NOT NULL,
+    PRIMARY KEY (sample_id, kind, contributor_id)
+);
+
+CREATE INDEX idx_usage_samples_occurred_at ON usage_samples(occurred_at);
+CREATE INDEX idx_usage_samples_thread_occurred_at ON usage_samples(thread_id, occurred_at);
+CREATE INDEX idx_usage_sample_contributors_kind ON usage_sample_contributors(kind, contributor_id);

codex-rs/state/src/lib.rs

@@ -48,6 +48,11 @@ pub use model::ThreadGoalStatus;
 pub use model::ThreadMetadata;
 pub use model::ThreadMetadataBuilder;
 pub use model::ThreadsPage;
+pub use model::UsageEntry;
+pub use model::UsageHeadline;
+pub use model::UsageRange;
+pub use model::UsageReport;
+pub use model::UsageSample;
 pub use runtime::GoalAccountingMode;
 pub use runtime::GoalAccountingOutcome;
 pub use runtime::GoalStore;

codex-rs/state/src/model/mod.rs

@@ -5,6 +5,7 @@ mod log;
 mod memories;
 mod thread_goal;
 mod thread_metadata;
+mod usage;
 
 pub use agent_job::AgentJob;
 pub use agent_job::AgentJobCreateParams;
@@ -34,6 +35,11 @@ pub use thread_metadata::SortKey;
 pub use thread_metadata::ThreadMetadata;
 pub use thread_metadata::ThreadMetadataBuilder;
 pub use thread_metadata::ThreadsPage;
+pub use usage::UsageEntry;
+pub use usage::UsageHeadline;
+pub use usage::UsageRange;
+pub use usage::UsageReport;
+pub use usage::UsageSample;
 
 pub(crate) use agent_job::AgentJobItemRow;
 pub(crate) use agent_job::AgentJobRow;

codex-rs/state/src/model/usage.rs

@@ -0,0 +1,53 @@
+use codex_protocol::protocol::UsageAttributionItem;
+use codex_protocol::protocol::UsageContributorKind;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum UsageRange {
+    Day,
+    Week,
+}
+
+impl UsageRange {
+    pub(crate) fn seconds(self) -> i64 {
+        match self {
+            Self::Day => 24 * 60 * 60,
+            Self::Week => 7 * 24 * 60 * 60,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct UsageEntry {
+    pub kind: UsageContributorKind,
+    pub id: String,
+    pub label: String,
+    pub attributed_tokens: i64,
+    pub percent_of_usage: u8,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct UsageHeadline {
+    pub entry: UsageEntry,
+    pub note: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct UsageReport {
+    pub range: UsageRange,
+    pub generated_at: i64,
+    pub tracked_from: Option<i64>,
+    pub total_tokens: i64,
+    pub headline: Option<UsageHeadline>,
+    pub skills: Vec<UsageEntry>,
+    pub subagents: Vec<UsageEntry>,
+    pub agent_tasks: Vec<UsageEntry>,
+    pub apps: Vec<UsageEntry>,
+    pub mcp_servers: Vec<UsageEntry>,
+    pub plugins: Vec<UsageEntry>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct UsageSample {
+    pub thread_id: codex_protocol::ThreadId,
+    pub attribution: UsageAttributionItem,
+}

codex-rs/state/src/runtime.rs

@@ -65,6 +65,7 @@ mod remote_control;
 #[cfg(test)]
 mod test_support;
 mod threads;
+mod usage;
 
 pub use goals::GoalAccountingMode;
 pub use goals::GoalAccountingOutcome;
@@ -237,6 +238,12 @@ impl StateRuntime {
                 logs_path.display(),
             );
         }
+        if let Err(err) = runtime.run_usage_startup_maintenance().await {
+            warn!(
+                "failed to run startup maintenance for usage data in state db at {}: {err}",
+                state_path.display(),
+            );
+        }
         Ok(runtime)
     }