mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
andrewgu-oai
c62e9df39f
## Why Resuming or forking a thread loses the session flag indicating that reported token usage already includes retained reasoning. Counting that reasoning again can trigger automatic compaction too early. ## What changed Persist `last_usage_includes_reasoning` with `TokenUsageInfo` and use it when estimating context usage. Record the flag with completed sampling and compaction usage, and mark local estimates as including reasoning. Older records without the field retain the existing accounting behavior. ## Testing Add coverage for saved inclusive, exclusive, and legacy token counts. Extend the encrypted-reasoning compaction test to verify the compaction threshold for live, forked, and cold-resumed threads. GitOrigin-RevId: 9f22262a693111ec0406dce11804eb4844ad8983
421 lines
12 KiB
Rust
421 lines
12 KiB
Rust
use anyhow::Result;
|
|
use codex_protocol::SessionId;
|
|
use codex_protocol::ThreadId;
|
|
use codex_protocol::protocol::EventMsg;
|
|
use codex_protocol::protocol::GitInfo;
|
|
use codex_protocol::protocol::SessionMeta;
|
|
use codex_protocol::protocol::SessionMetaLine;
|
|
use codex_protocol::protocol::SessionSource;
|
|
use codex_protocol::protocol::ThreadHistoryMode;
|
|
use codex_protocol::protocol::ThreadSource;
|
|
use codex_protocol::protocol::TokenCountEvent;
|
|
use codex_protocol::protocol::TokenUsage;
|
|
use codex_protocol::protocol::TokenUsageInfo;
|
|
use core_test_support::test_path_buf;
|
|
use serde_json::json;
|
|
use std::fs;
|
|
use std::fs::FileTimes;
|
|
use std::path::Path;
|
|
use std::path::PathBuf;
|
|
use uuid::Uuid;
|
|
|
|
pub fn rollout_path(codex_home: &Path, filename_ts: &str, thread_id: &str) -> PathBuf {
|
|
let year = &filename_ts[0..4];
|
|
let month = &filename_ts[5..7];
|
|
let day = &filename_ts[8..10];
|
|
codex_home
|
|
.join("sessions")
|
|
.join(year)
|
|
.join(month)
|
|
.join(day)
|
|
.join(format!("rollout-{filename_ts}-{thread_id}.jsonl"))
|
|
}
|
|
|
|
/// Create a minimal rollout file under `CODEX_HOME/sessions/YYYY/MM/DD/`.
|
|
///
|
|
/// - `filename_ts` is the filename timestamp component in `YYYY-MM-DDThh-mm-ss` format.
|
|
/// - `meta_rfc3339` is the envelope timestamp used in JSON lines.
|
|
/// - `preview` is the user message preview text.
|
|
/// - `model_provider` optionally sets the provider in the session meta payload.
|
|
///
|
|
/// Returns the generated conversation/session UUID as a string.
|
|
pub fn create_fake_rollout(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
model_provider: Option<&str>,
|
|
git_info: Option<GitInfo>,
|
|
) -> Result<String> {
|
|
create_fake_rollout_with_source(
|
|
codex_home,
|
|
filename_ts,
|
|
meta_rfc3339,
|
|
preview,
|
|
model_provider,
|
|
git_info,
|
|
SessionSource::Cli,
|
|
)
|
|
}
|
|
|
|
/// Creates a minimal paginated rollout with ordinalized JSONL records.
|
|
pub fn create_fake_paginated_rollout(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
model_provider: Option<&str>,
|
|
git_info: Option<GitInfo>,
|
|
) -> Result<String> {
|
|
let thread_id = create_fake_rollout(
|
|
codex_home,
|
|
filename_ts,
|
|
meta_rfc3339,
|
|
preview,
|
|
model_provider,
|
|
git_info,
|
|
)?;
|
|
let path = rollout_path(codex_home, filename_ts, &thread_id);
|
|
let mut lines = fs::read_to_string(path.as_path())?
|
|
.lines()
|
|
.map(serde_json::from_str::<serde_json::Value>)
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
lines[0]["payload"]["history_mode"] = serde_json::to_value(ThreadHistoryMode::Paginated)?;
|
|
for (ordinal, line) in lines.iter_mut().enumerate() {
|
|
line["ordinal"] = serde_json::to_value(ordinal)?;
|
|
}
|
|
let contents = lines
|
|
.into_iter()
|
|
.map(|line| line.to_string())
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
fs::write(path, format!("{contents}\n"))?;
|
|
Ok(thread_id)
|
|
}
|
|
|
|
/// Creates a minimal rollout whose history includes a persisted token usage event.
|
|
///
|
|
/// Resume and fork tests use this fixture to verify lifecycle replay of restored
|
|
/// usage without starting a model turn. The exact token values are intentionally
|
|
/// non-zero and asymmetric so assertions catch swapped total/last fields and
|
|
/// dropped cached or reasoning counters.
|
|
pub fn create_fake_rollout_with_token_usage(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
model_provider: Option<&str>,
|
|
) -> Result<String> {
|
|
let thread_id = create_fake_rollout(
|
|
codex_home,
|
|
filename_ts,
|
|
meta_rfc3339,
|
|
preview,
|
|
model_provider,
|
|
/*git_info*/ None,
|
|
)?;
|
|
let payload = serde_json::to_value(EventMsg::TokenCount(TokenCountEvent {
|
|
info: Some(TokenUsageInfo {
|
|
total_token_usage: TokenUsage {
|
|
input_tokens: 120,
|
|
cached_input_tokens: 20,
|
|
cache_write_input_tokens: 0,
|
|
output_tokens: 30,
|
|
reasoning_output_tokens: 10,
|
|
total_tokens: 150,
|
|
codex_rollout_budget_units: None,
|
|
},
|
|
last_usage_includes_reasoning: None,
|
|
last_token_usage: TokenUsage {
|
|
input_tokens: 70,
|
|
cached_input_tokens: 10,
|
|
cache_write_input_tokens: 0,
|
|
output_tokens: 20,
|
|
reasoning_output_tokens: 5,
|
|
total_tokens: 90,
|
|
codex_rollout_budget_units: None,
|
|
},
|
|
model_context_window: Some(200_000),
|
|
}),
|
|
rate_limits: None,
|
|
}))?;
|
|
let file_path = rollout_path(codex_home, filename_ts, &thread_id);
|
|
let line = json!({
|
|
"timestamp": meta_rfc3339,
|
|
"type": "event_msg",
|
|
"payload": payload
|
|
})
|
|
.to_string();
|
|
fs::write(
|
|
&file_path,
|
|
format!("{}{}\n", fs::read_to_string(&file_path)?, line),
|
|
)?;
|
|
Ok(thread_id)
|
|
}
|
|
|
|
/// Create a minimal rollout file with an explicit session source.
|
|
pub fn create_fake_rollout_with_source(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
model_provider: Option<&str>,
|
|
git_info: Option<GitInfo>,
|
|
source: SessionSource,
|
|
) -> Result<String> {
|
|
create_fake_rollout_with_session_and_thread_source(
|
|
codex_home,
|
|
filename_ts,
|
|
meta_rfc3339,
|
|
preview,
|
|
model_provider,
|
|
git_info,
|
|
source,
|
|
/*thread_source*/ None,
|
|
)
|
|
}
|
|
|
|
/// Create a minimal rollout file with explicit session and thread sources.
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn create_fake_rollout_with_session_and_thread_source(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
model_provider: Option<&str>,
|
|
git_info: Option<GitInfo>,
|
|
source: SessionSource,
|
|
thread_source: Option<ThreadSource>,
|
|
) -> Result<String> {
|
|
create_fake_rollout_with_source_and_parent_thread_id(
|
|
codex_home,
|
|
filename_ts,
|
|
meta_rfc3339,
|
|
preview,
|
|
model_provider,
|
|
git_info,
|
|
source,
|
|
thread_source,
|
|
/*session_id*/ None,
|
|
/*parent_thread_id*/ None,
|
|
)
|
|
}
|
|
|
|
/// Create a minimal rollout file with an explicit root session and control parent.
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn create_fake_parented_rollout_with_source(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
model_provider: Option<&str>,
|
|
git_info: Option<GitInfo>,
|
|
source: SessionSource,
|
|
session_id: SessionId,
|
|
parent_thread_id: ThreadId,
|
|
) -> Result<String> {
|
|
create_fake_rollout_with_source_and_parent_thread_id(
|
|
codex_home,
|
|
filename_ts,
|
|
meta_rfc3339,
|
|
preview,
|
|
model_provider,
|
|
git_info,
|
|
source,
|
|
/*thread_source*/ None,
|
|
Some(session_id),
|
|
Some(parent_thread_id),
|
|
)
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn create_fake_rollout_with_source_and_parent_thread_id(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
model_provider: Option<&str>,
|
|
git_info: Option<GitInfo>,
|
|
source: SessionSource,
|
|
thread_source: Option<ThreadSource>,
|
|
session_id: Option<SessionId>,
|
|
parent_thread_id: Option<ThreadId>,
|
|
) -> Result<String> {
|
|
let uuid = Uuid::new_v4();
|
|
let uuid_str = uuid.to_string();
|
|
let conversation_id = ThreadId::from_string(&uuid_str)?;
|
|
let session_id = session_id.unwrap_or_else(|| conversation_id.into());
|
|
|
|
let file_path = rollout_path(codex_home, filename_ts, &uuid_str);
|
|
let dir = file_path
|
|
.parent()
|
|
.ok_or_else(|| anyhow::anyhow!("missing rollout parent directory"))?;
|
|
fs::create_dir_all(dir)?;
|
|
|
|
// Build JSONL lines
|
|
let meta = SessionMeta {
|
|
attached_thread: false,
|
|
creator_user_id: None,
|
|
creator_account_id: None,
|
|
session_id,
|
|
id: conversation_id,
|
|
forked_from_id: None,
|
|
forked_from_ordinal_exclusive: None,
|
|
parent_thread_id,
|
|
timestamp: meta_rfc3339.to_string(),
|
|
cwd: test_path_buf("/"),
|
|
runtime_workspace_roots: None,
|
|
originator: "codex".to_string(),
|
|
cli_version: "0.0.0".to_string(),
|
|
source,
|
|
thread_source,
|
|
agent_path: None,
|
|
agent_nickname: None,
|
|
agent_role: None,
|
|
model_provider: model_provider.map(str::to_string),
|
|
base_instructions: None,
|
|
dynamic_tools: None,
|
|
selected_capability_roots: Vec::new(),
|
|
memory_mode: None,
|
|
history_mode: Default::default(),
|
|
history_base: None,
|
|
subagent_history_start_ordinal: None,
|
|
multi_agent_version: None,
|
|
context_window: None,
|
|
};
|
|
let payload = serde_json::to_value(SessionMetaLine {
|
|
meta,
|
|
git: git_info,
|
|
})?;
|
|
|
|
let lines = [
|
|
json!({
|
|
"timestamp": meta_rfc3339,
|
|
"type": "session_meta",
|
|
"payload": payload
|
|
})
|
|
.to_string(),
|
|
json!({
|
|
"timestamp": meta_rfc3339,
|
|
"type":"response_item",
|
|
"payload": {
|
|
"type":"message",
|
|
"role":"user",
|
|
"content":[{"type":"input_text","text": preview}]
|
|
}
|
|
})
|
|
.to_string(),
|
|
json!({
|
|
"timestamp": meta_rfc3339,
|
|
"type":"event_msg",
|
|
"payload": {
|
|
"type":"user_message",
|
|
"message": preview,
|
|
"kind": "plain"
|
|
}
|
|
})
|
|
.to_string(),
|
|
];
|
|
|
|
fs::write(&file_path, lines.join("\n") + "\n")?;
|
|
let parsed = chrono::DateTime::parse_from_rfc3339(meta_rfc3339)?.with_timezone(&chrono::Utc);
|
|
let times = FileTimes::new().set_modified(parsed.into());
|
|
std::fs::OpenOptions::new()
|
|
.append(true)
|
|
.open(&file_path)?
|
|
.set_times(times)?;
|
|
Ok(uuid_str)
|
|
}
|
|
|
|
pub fn create_fake_rollout_with_text_elements(
|
|
codex_home: &Path,
|
|
filename_ts: &str,
|
|
meta_rfc3339: &str,
|
|
preview: &str,
|
|
text_elements: Vec<serde_json::Value>,
|
|
model_provider: Option<&str>,
|
|
git_info: Option<GitInfo>,
|
|
) -> Result<String> {
|
|
let uuid = Uuid::new_v4();
|
|
let uuid_str = uuid.to_string();
|
|
let conversation_id = ThreadId::from_string(&uuid_str)?;
|
|
|
|
// sessions/YYYY/MM/DD derived from filename_ts (YYYY-MM-DDThh-mm-ss)
|
|
let year = &filename_ts[0..4];
|
|
let month = &filename_ts[5..7];
|
|
let day = &filename_ts[8..10];
|
|
let dir = codex_home.join("sessions").join(year).join(month).join(day);
|
|
fs::create_dir_all(&dir)?;
|
|
|
|
let file_path = dir.join(format!("rollout-{filename_ts}-{uuid}.jsonl"));
|
|
|
|
// Build JSONL lines
|
|
let meta = SessionMeta {
|
|
attached_thread: false,
|
|
creator_user_id: None,
|
|
creator_account_id: None,
|
|
session_id: conversation_id.into(),
|
|
id: conversation_id,
|
|
forked_from_id: None,
|
|
forked_from_ordinal_exclusive: None,
|
|
parent_thread_id: None,
|
|
timestamp: meta_rfc3339.to_string(),
|
|
cwd: test_path_buf("/"),
|
|
runtime_workspace_roots: None,
|
|
originator: "codex".to_string(),
|
|
cli_version: "0.0.0".to_string(),
|
|
source: SessionSource::Cli,
|
|
thread_source: None,
|
|
agent_path: None,
|
|
agent_nickname: None,
|
|
agent_role: None,
|
|
model_provider: model_provider.map(str::to_string),
|
|
base_instructions: None,
|
|
dynamic_tools: None,
|
|
selected_capability_roots: Vec::new(),
|
|
memory_mode: None,
|
|
history_mode: Default::default(),
|
|
history_base: None,
|
|
subagent_history_start_ordinal: None,
|
|
multi_agent_version: None,
|
|
context_window: None,
|
|
};
|
|
let payload = serde_json::to_value(SessionMetaLine {
|
|
meta,
|
|
git: git_info,
|
|
})?;
|
|
|
|
let lines = [
|
|
json!( {
|
|
"timestamp": meta_rfc3339,
|
|
"type": "session_meta",
|
|
"payload": payload
|
|
})
|
|
.to_string(),
|
|
json!( {
|
|
"timestamp": meta_rfc3339,
|
|
"type":"response_item",
|
|
"payload": {
|
|
"type":"message",
|
|
"role":"user",
|
|
"content":[{"type":"input_text","text": preview}]
|
|
}
|
|
})
|
|
.to_string(),
|
|
json!( {
|
|
"timestamp": meta_rfc3339,
|
|
"type":"event_msg",
|
|
"payload": {
|
|
"type":"user_message",
|
|
"message": preview,
|
|
"text_elements": text_elements,
|
|
"local_images": []
|
|
}
|
|
})
|
|
.to_string(),
|
|
];
|
|
|
|
fs::write(file_path, lines.join("\n") + "\n")?;
|
|
Ok(uuid_str)
|
|
}
|