diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index dfae092ab7..2e2ec60434 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -5400,6 +5400,19 @@ dependencies = [ "toml 0.9.11+spec-1.1.0", ] +[[package]] +name = "codex-utils-model-facing-id" +version = "0.0.0" +dependencies = [ + "blake3", + "pretty_assertions", + "schemars 0.8.22", + "serde", + "serde_json", + "thiserror 2.0.18", + "tracing", +] + [[package]] name = "codex-utils-oss" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index cbd41d41d2..0acadbfe5c 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -147,6 +147,7 @@ members = [ "utils/cli", "utils/elapsed", "utils/sandbox-summary", + "utils/model-facing-id", "utils/sleep-inhibitor", "utils/approval-presets", "utils/oss", diff --git a/codex-rs/utils/model-facing-id/BUILD.bazel b/codex-rs/utils/model-facing-id/BUILD.bazel new file mode 100644 index 0000000000..8e195c83f6 --- /dev/null +++ b/codex-rs/utils/model-facing-id/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "model-facing-id", + crate_name = "codex_utils_model_facing_id", +) diff --git a/codex-rs/utils/model-facing-id/Cargo.toml b/codex-rs/utils/model-facing-id/Cargo.toml new file mode 100644 index 0000000000..587da05235 --- /dev/null +++ b/codex-rs/utils/model-facing-id/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "codex-utils-model-facing-id" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lints] +workspace = true + +[lib] +doctest = false +test = false + +[dependencies] +blake3.workspace = true +schemars.workspace = true +serde = { workspace = true, features = ["derive"] } +thiserror.workspace = true +tracing.workspace = true + +[dev-dependencies] +serde_json.workspace = true +pretty_assertions.workspace = true diff --git a/codex-rs/utils/model-facing-id/src/lib.rs b/codex-rs/utils/model-facing-id/src/lib.rs new file mode 100644 index 0000000000..c93dbd2f48 --- /dev/null +++ b/codex-rs/utils/model-facing-id/src/lib.rs @@ -0,0 +1,197 @@ +//! Issues a salt followed by shuffled decimal IDs from `000`–`999`, then `001000`–`999999`. +//! Each pool uses a four-round Feistel permutation with BLAKE3 and cycle-walking. +//! The seed can be varied across sessions to change the order in which IDs are issued. + +use std::borrow::Cow; +use std::fmt; +use std::ops::RangeInclusive; + +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use thiserror::Error; + +const VERSION: u32 = 1; +// Our tokenizers encode every decimal ID of up to three digits as a single token. +const GROUP_SIZE: u64 = 1_000; + +/// Stores the salt and numeric value separately. Display pads the value to groups of three digits. +/// With salt `environment`, zero renders as `environment000` and 1000 as `environment001000`. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct ModelFacingId { + salt: Cow<'static, str>, + value: u64, +} + +impl fmt::Display for ModelFacingId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}{}", self.salt, self.digits()) + } +} + +impl ModelFacingId { + /// Formats the padded digits for protocols that supply their own type prefix. + pub fn digits(&self) -> impl fmt::Display + '_ { + fmt::from_fn(|f| { + write!( + f, + "{:0width$}", + self.value, + width = decimal_width(self.value) + ) + }) + } +} + +fn decimal_width(value: u64) -> usize { + (value.checked_ilog10().unwrap_or(0) as usize + 1).next_multiple_of(3) +} + +/// Generates token-efficient identifiers for model context. +#[derive(Debug)] +pub struct ModelFacingIdGenerator { + state: ModelFacingIdGeneratorState, +} + +impl ModelFacingIdGenerator { + /// Starts a sequence for one ID kind using a shared session seed. + /// The salt also prefixes each displayed ID. Use the [tokenizer](https://platform.openai.com/tokenizer) + /// to choose a stable, single-token salt such as `environment`. + /// + /// # Panics + /// + /// Panics if the salt is empty. + pub fn new(seed: u64, salt: &'static str) -> Self { + assert!(!salt.is_empty(), "a model ID salt must be nonempty"); + let mut hasher = blake3::Hasher::new(); + hasher.update(&seed.to_le_bytes()); + hasher.update(salt.as_bytes()); + let mut seed_bytes = [0; 8]; + hasher.finalize_xof().fill(&mut seed_bytes); + let derived_seed = u64::from_le_bytes(seed_bytes); + Self { + state: ModelFacingIdGeneratorState::new( + VERSION, + derived_seed, + salt, + /*next_ordinal*/ 0, + ), + } + } + + /// Issues the next ID. After all `u64` values are issued, the sequence repeats from zero. + pub fn fetch_next(&mut self) -> ModelFacingId { + let ordinal = self.state.next_ordinal; + let mut pool = 0..=GROUP_SIZE - 1; + while ordinal > *pool.end() { + let start = pool.end() + 1; + // The final decimal pool ends at u64::MAX, which has no representable successor. + let last = start + .checked_mul(GROUP_SIZE) + .map_or(u64::MAX, |end| end - 1); + pool = start..=last; + } + let id = permute(self.state.seed, pool, ordinal); + let (next, wrapped) = ordinal.overflowing_add(1); + if wrapped { + tracing::error!( + "model-facing ID generator exhausted all u64 values and wrapped to zero" + ); + } + self.state.next_ordinal = next; + ModelFacingId { + salt: self.state.salt.clone(), + value: id, + } + } + + /// Captures the position of the next allocation. + pub fn snapshot(&self) -> ModelFacingIdGeneratorState { + self.state.clone() + } + + /// Continues the snapshot's sequence from its next unissued ordinal. + pub fn restore(state: ModelFacingIdGeneratorState) -> Result { + if state.version != VERSION { + return Err(RestoreError::UnsupportedVersion(state.version)); + } + if state.salt.is_empty() { + return Err(RestoreError::InvalidSalt); + } + Ok(Self { state }) + } +} + +/// Persisted sequence version, derived seed, salt, and zero-based next allocation ordinal. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +pub struct ModelFacingIdGeneratorState { + version: u32, + seed: u64, + salt: Cow<'static, str>, + next_ordinal: u64, +} + +impl ModelFacingIdGeneratorState { + /// Returns the namespace this snapshot belongs to. + pub fn salt(&self) -> &str { + &self.salt + } + + /// Rebuilds a snapshot from its stored fields. [`ModelFacingIdGenerator::restore`] validates it. + pub fn new( + version: u32, + seed: u64, + salt: impl Into>, + next_ordinal: u64, + ) -> Self { + Self { + version, + seed, + salt: salt.into(), + next_ordinal, + } + } +} + +/// The snapshot cannot be used to resume an allocation sequence. +#[derive(Debug, PartialEq, Eq, Error)] +pub enum RestoreError { + #[error("unsupported model-facing ID generator version {0}")] + UnsupportedVersion(u32), + #[error("a model ID salt must be nonempty")] + InvalidSalt, +} + +// Use the smallest balanced Feistel domain that covers the pool to limit cycle-walking. +// Both the four rounds and their input encoding are part of the snapshot format. +fn permute(seed: u64, pool: RangeInclusive, ordinal: u64) -> u64 { + let start = *pool.start(); + let size = pool.end() - start + 1; + let bits = u64::BITS - (size - 1).leading_zeros(); + let half_bits = bits.div_ceil(2); + let half_mask = (1_u64 << half_bits) - 1; + // The seed, pool index, and round occupy 17 bytes. The pool width fixes the right half's length. + let input_len = 17 + half_bits.div_ceil(8) as usize; + let mut input = [0; 25]; + input[..8].copy_from_slice(&seed.to_le_bytes()); + input[8..16].copy_from_slice(&(start / GROUP_SIZE).to_le_bytes()); + + let mut slot = ordinal - start; + loop { + let (mut left, mut right) = (slot >> half_bits, slot & half_mask); + for round in 0..4 { + input[16] = round; + input[17..].copy_from_slice(&right.to_le_bytes()); + let hash = blake3::hash(&input[..input_len]); + let mut bytes = [0; 8]; + bytes.copy_from_slice(&hash.as_bytes()[..8]); + let f = u64::from_le_bytes(bytes) & half_mask; + (left, right) = (right, left ^ f); + } + slot = (left << half_bits) | right; + // Following the permutation cycle must eventually return to the starting slot. + if slot < size { + return start + slot; + } + } +} diff --git a/codex-rs/utils/model-facing-id/tests/generator.rs b/codex-rs/utils/model-facing-id/tests/generator.rs new file mode 100644 index 0000000000..805c6e9461 --- /dev/null +++ b/codex-rs/utils/model-facing-id/tests/generator.rs @@ -0,0 +1,193 @@ +//! Checks pool boundaries, snapshot restore, and the v1 sequence through the public API. + +use std::collections::BTreeSet; +use std::error::Error; + +use codex_utils_model_facing_id::ModelFacingId; +use codex_utils_model_facing_id::ModelFacingIdGenerator; +use codex_utils_model_facing_id::ModelFacingIdGeneratorState; +use codex_utils_model_facing_id::RestoreError; +use pretty_assertions::assert_eq; +use serde_json::json; + +/// Every three-digit suffix is issued before the first ID drawn from the six-digit pool. +#[test] +fn exhausts_three_digit_ids_before_opening_the_six_digit_pool() { + let mut generator = ModelFacingIdGenerator::new(/*seed*/ 42, "environment"); + let mut actual = BTreeSet::new(); + for _ in 0..1_000 { + actual.insert(generator.fetch_next().to_string()); + } + let expected: BTreeSet<_> = (0..1_000).map(|id| format!("environment{id:03}")).collect(); + assert_eq!(actual, expected); + assert_eq!(generator.fetch_next().to_string(), "environment825109"); +} + +/// IDs can have the same display text and still round-trip with distinct salts and values. +#[test] +fn persisted_ids_keep_the_salt_and_value_separate() -> Result<(), serde_json::Error> { + for (salt, value) in [("env123", 0), ("env", 123_000)] { + let saved = json!({ "salt": salt, "value": value }); + let id: ModelFacingId = serde_json::from_value(saved.clone())?; + assert_eq!(id.to_string(), "env123000"); + assert_eq!(serde_json::to_value(id)?, saved); + } + Ok(()) +} + +/// A saved cursor resumes the uninterrupted sequence both inside and between pools. +#[test] +fn serialized_snapshots_resume_the_sequence() -> Result<(), Box> { + let derived_seed = 0x9e6e_7c74_2945_897e; + for next_ordinal in [437, 1_000, 1_437] { + let mut original = ModelFacingIdGenerator::new(u64::MAX, "env2"); + for _ in 0..next_ordinal { + original.fetch_next(); + } + let snapshot = original.snapshot(); + assert_eq!( + ModelFacingIdGeneratorState::new( + /*version*/ 1, + derived_seed, + "env2", + next_ordinal + ), + snapshot + ); + let saved = serde_json::to_value(snapshot)?; + assert_eq!( + saved, + json!({"version": 1, "seed": derived_seed, "salt": "env2", "next_ordinal": next_ordinal}) + ); + let mut resumed = ModelFacingIdGenerator::restore(serde_json::from_value(saved)?)?; + for _ in 0..1_001 { + assert_eq!(resumed.fetch_next(), original.fetch_next()); + } + } + Ok(()) +} + +/// V1 must keep the initial outputs and repeated cycle walks stable. +#[test] +fn v1_known_outputs() -> Result<(), Box> { + let seed = 0x1234_5678_9abc_def0; + let state = ModelFacingIdGeneratorState::new( + /*version*/ 1, + seed, + "environment", + /*next_ordinal*/ 0, + ); + let mut generator = ModelFacingIdGenerator::restore(state)?; + for id in ["895", "930", "155", "141", "916", "912", "678", "430"] { + assert_eq!( + generator.fetch_next().to_string(), + format!("environment{id}") + ); + } + // Ordinal 1,605 visits two excluded slots before returning to the six-digit pool. + let state = ModelFacingIdGeneratorState::new( + /*version*/ 1, + seed, + "environment", + /*next_ordinal*/ 1_605, + ); + let mut resumed = ModelFacingIdGenerator::restore(state)?; + assert_eq!(resumed.fetch_next().to_string(), "environment765577"); + Ok(()) +} + +/// The last ID at each width and the first at the next width stay stable across restores. +#[test] +fn v1_known_outputs_at_every_pool_boundary() -> Result<(), Box> { + let seed = 0x1234_5678_9abc_def0; + for (boundary, last, first) in [ + (1_000, "553", "661326"), + (1_000_000, "753874", "286136338"), + (1_000_000_000, "839882170", "813648058018"), + (1_000_000_000_000, "634462456979", "134726697070558"), + ( + 1_000_000_000_000_000, + "326088534089890", + "563366069705075074", + ), + ( + 1_000_000_000_000_000_000, + "633936232025849076", + "012827849051678362461", + ), + ] { + let state = ModelFacingIdGeneratorState::new(/*version*/ 1, seed, "env", boundary - 1); + let mut generator = ModelFacingIdGenerator::restore(state)?; + assert_eq!(generator.fetch_next().to_string(), format!("env{last}")); + let mut resumed = ModelFacingIdGenerator::restore(generator.snapshot())?; + let next = generator.fetch_next(); + assert_eq!(next.to_string(), format!("env{first}")); + assert_eq!(resumed.fetch_next(), next); + } + Ok(()) +} + +/// The same session seed can drive several ID kinds, each with its own resumable seed. +#[test] +fn salts_derive_the_seed_stored_in_snapshots() { + for (seed, salt, derived_seed) in [ + (42, "environment", 0xc9b5_d0ac_9722_ea42), + (42, "context-window", 0x2747_db8a_6513_485b), + (43, "environment", 0x70aa_a9fd_4bf2_6e41), + ] { + assert_eq!( + ModelFacingIdGenerator::new(seed, salt).snapshot(), + ModelFacingIdGeneratorState::new( + /*version*/ 1, + derived_seed, + salt, + /*next_ordinal*/ 0 + ) + ); + } +} + +/// The final ordinal belongs to the last decimal pool. Its successor restarts the sequence. +#[test] +fn wraps_after_the_last_u64_ordinal() -> Result<(), Box> { + let seed = 17; + let last = ModelFacingIdGeneratorState::new(/*version*/ 1, seed, "env", u64::MAX); + let mut generator = ModelFacingIdGenerator::restore(last)?; + let text = generator.fetch_next().to_string(); + let last_id = text + .strip_prefix("env") + .expect("restored salt") + .parse::()?; + assert!(last_id >= 1_000_000_000_000_000_000); + + let first = + ModelFacingIdGeneratorState::new(/*version*/ 1, seed, "env", /*next_ordinal*/ 0); + assert_eq!(generator.snapshot(), first); + let mut fresh = ModelFacingIdGenerator::restore(first)?; + let mut resumed = ModelFacingIdGenerator::restore(generator.snapshot())?; + assert_eq!(resumed.fetch_next(), fresh.fetch_next()); + assert_eq!(resumed.snapshot(), fresh.snapshot()); + Ok(()) +} + +#[test] +fn rejects_unsupported_snapshot_versions() { + let saved = ModelFacingIdGeneratorState::new( + /*version*/ 2, /*seed*/ 0, "env", /*next_ordinal*/ 0, + ); + assert_eq!( + ModelFacingIdGenerator::restore(saved).unwrap_err(), + RestoreError::UnsupportedVersion(2) + ); +} + +#[test] +fn rejects_empty_salts_on_restore() { + let state = ModelFacingIdGeneratorState::new( + /*version*/ 1, /*seed*/ 42, "", /*next_ordinal*/ 0, + ); + assert_eq!( + ModelFacingIdGenerator::restore(state).expect_err("empty salt"), + RestoreError::InvalidSalt + ); +}