mirror of
https://github.com/openai/codex.git
synced 2026-09-05 15:18:41 +00:00
feat(usage): add local usage storage
This commit is contained in:
@@ -1925,6 +1925,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,
|
||||
|
||||
@@ -216,6 +216,66 @@ async fn load_rollout_items_skips_legacy_ghost_snapshot_lines() -> std::io::Resu
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_rollout_items_skips_legacy_usage_attribution_lines() -> std::io::Result<()> {
|
||||
let home = TempDir::new().expect("temp dir");
|
||||
let rollout_path = home.path().join("rollout.jsonl");
|
||||
let mut file = File::create(&rollout_path)?;
|
||||
let thread_id = ThreadId::new();
|
||||
let ts = "2025-01-03T12:00:00Z";
|
||||
|
||||
writeln!(
|
||||
file,
|
||||
"{}",
|
||||
serde_json::json!({
|
||||
"timestamp": ts,
|
||||
"type": "session_meta",
|
||||
"payload": {
|
||||
"id": thread_id,
|
||||
"timestamp": ts,
|
||||
"cwd": ".",
|
||||
"originator": "test_originator",
|
||||
"cli_version": "test_version",
|
||||
"source": "cli",
|
||||
"model_provider": "test-provider",
|
||||
},
|
||||
})
|
||||
)?;
|
||||
writeln!(
|
||||
file,
|
||||
"{}",
|
||||
serde_json::json!({
|
||||
"timestamp": ts,
|
||||
"type": "usage_attribution",
|
||||
"payload": {
|
||||
"sample_id": "sample",
|
||||
"turn_id": "turn",
|
||||
"response_id": "response",
|
||||
"occurred_at": 1_700_000_000,
|
||||
"token_usage": {
|
||||
"input_tokens": 1,
|
||||
"cached_input_tokens": 0,
|
||||
"output_tokens": 1,
|
||||
"reasoning_output_tokens": 0,
|
||||
"total_tokens": 2,
|
||||
},
|
||||
"prompt_estimated_tokens": 1,
|
||||
"contributors": [],
|
||||
},
|
||||
})
|
||||
)?;
|
||||
|
||||
let (items, loaded_thread_id, parse_errors) =
|
||||
RolloutRecorder::load_rollout_items(&rollout_path).await?;
|
||||
|
||||
assert_eq!(loaded_thread_id, Some(thread_id));
|
||||
assert_eq!(parse_errors, 1);
|
||||
assert_eq!(items.len(), 1);
|
||||
assert!(matches!(items[0], RolloutItem::SessionMeta(_)));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_rollout_items_preserves_legacy_guardian_assessment_lines() -> std::io::Result<()> {
|
||||
let home = TempDir::new().expect("temp dir");
|
||||
|
||||
29
codex-rs/state/migrations/0035_usage_samples.sql
Normal file
29
codex-rs/state/migrations/0035_usage_samples.sql
Normal file
@@ -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);
|
||||
@@ -51,6 +51,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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
53
codex-rs/state/src/model/usage.rs
Normal file
53
codex-rs/state/src/model/usage.rs
Normal file
@@ -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,
|
||||
}
|
||||
@@ -66,6 +66,7 @@ mod remote_control;
|
||||
#[cfg(test)]
|
||||
mod test_support;
|
||||
mod threads;
|
||||
mod usage;
|
||||
|
||||
pub use goals::GoalAccountingMode;
|
||||
pub use goals::GoalAccountingOutcome;
|
||||
@@ -267,6 +268,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)
|
||||
}
|
||||
|
||||
|
||||
1184
codex-rs/state/src/runtime/usage.rs
Normal file
1184
codex-rs/state/src/runtime/usage.rs
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user