Replace the footer's placeholder usage with a real observation ledger
The footer showed invented percentages. It now shows what two harnesses actually report, or says it does not know. lumbridge-core gains an append-only per-profile UsageLedger and a projection that labels every derived value estimated, withholds a burn rate from a single sample, withholds a window fraction with no reported ceiling, withholds an exhaustion estimate that lands after the reset, and reports an expired window as rolled over rather than freezing its last percentage. A missing fact renders as missing, never as zero. (0012) lumbridge-harness is the impure side: processes, clocks, and untrusted wire text in, observations out. Three adapters: - Codex's account/rateLimits/read over the app-server's JSON-RPC stdio. The client cannot express a request outside a two-variant enum and answers every server-to-client request with -32601, so a harness asking Lumbridge for a credential is refused by construction. (0013) - Claude Code's session transcripts, as a byte-offset tail follower that reports nothing until the backlog is read to EOF — a partially-read backlog is indistinguishable from a burst of spend, and the first run against 20 MB reported forty-six billion tokens an hour. The parser models four counters, so the conversations in those files are not representable. (0014) - Claude Code's five-hour and seven-day subscription windows, via a bridge installed as its statusLine command. 0014 had claimed no such surface existed; it does, and the record is corrected in place rather than quietly edited. Lumbridge does not read the OAuth credential to call the account usage endpoint, which is what comparable tools do — AGENTS.md forbids it, and 0015 says so rather than leaving the gap unexplained. Also in here: a capability-check ordering fix in the workspace reducer, where the applied-request replay table was consulted before the capability check and so answered questions the caller had no right to ask; the GPUI spike wired to the live probes with per-harness gauges and provenance chips; and a launcher that matches its own window by PID, because GPUI sets WM_NAME but not _NET_WM_NAME and a title match never succeeded. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Metal Agent
co-authored by
Claude Opus 5
Claude Opus 5
parent
7fe84f71e2
commit
ef52aa7ce2
@@ -0,0 +1,324 @@
|
||||
//! A supervised, byte-capped stdio child.
|
||||
//!
|
||||
//! The child is launched with a cleared environment and a literal allowlist.
|
||||
//! It speaks a machine protocol, so it gets pipes rather than a PTY: a TTY
|
||||
//! would invite terminal control sequences in both directions for no benefit.
|
||||
//! Its stderr is discarded rather than captured, because a harness's diagnostic
|
||||
//! output is exactly the kind of text that ends up in a log carrying a path or
|
||||
//! a token fragment.
|
||||
//!
|
||||
//! Reading a pipe blocks. The child therefore splits into three owners: a
|
||||
//! reader that may block on its own thread, a writer, and a killer the
|
||||
//! supervising thread keeps. Shutdown kills first and joins second, so a child
|
||||
//! that has gone quiet — or hostile — can never hold the caller hostage.
|
||||
//!
|
||||
//! The child is put in its own process group and the group is what gets
|
||||
//! killed. A pipe reports end of file only when *every* write end closes, so
|
||||
//! killing the direct child alone is not enough: a launcher that execs the
|
||||
//! real program as a grandchild with inherited stdio — which is how the npm
|
||||
//! distribution of Codex works — leaves that grandchild holding our stdout,
|
||||
//! and the reader would block forever on a pipe nothing will ever write to.
|
||||
|
||||
use std::io::{BufRead, BufReader, Read, Write};
|
||||
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use crate::HarnessError;
|
||||
|
||||
/// Environment a probe child may inherit.
|
||||
///
|
||||
/// Literal names only: never a pattern and never a user-supplied name.
|
||||
///
|
||||
/// This deliberately mirrors but does not reuse the terminal crate's private
|
||||
/// allowlist. AGENTS.md keeps the harness and terminal boundaries separate,
|
||||
/// and the two lists genuinely differ: a probe needs `CODEX_HOME` so it reads
|
||||
/// the same account the user's panes use, and has no use for `SHELL`.
|
||||
/// `CODEX_HOME` is a filesystem path, never a secret.
|
||||
pub const PROBE_ENV_ALLOWLIST: &[&str] = &[
|
||||
"CODEX_HOME",
|
||||
"HOME",
|
||||
"LANG",
|
||||
"LC_ALL",
|
||||
"LC_CTYPE",
|
||||
"LOGNAME",
|
||||
"PATH",
|
||||
"TMPDIR",
|
||||
"USER",
|
||||
"XDG_CONFIG_HOME",
|
||||
"XDG_DATA_HOME",
|
||||
"XDG_RUNTIME_DIR",
|
||||
];
|
||||
|
||||
/// A single frame longer than this is treated as a fault rather than buffered.
|
||||
/// A hostile or broken relay must not be able to grow this process's memory.
|
||||
pub(crate) const MAX_LINE_BYTES: usize = 256 * 1024;
|
||||
|
||||
/// A freshly launched child, before its pipes are handed to their owners.
|
||||
pub(crate) struct ProbeChild {
|
||||
child: Arc<Mutex<Child>>,
|
||||
group: u32,
|
||||
stdin: ChildStdin,
|
||||
stdout: ChildStdout,
|
||||
}
|
||||
|
||||
impl ProbeChild {
|
||||
/// Launches a probe program with a minimal environment.
|
||||
///
|
||||
/// This runs on the caller's thread, following the same pattern as the
|
||||
/// runtime actor: start synchronously so a missing program is an immediate
|
||||
/// error, then transfer the pipes to their threads.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`HarnessError::EmptyProgram`] for a blank program name and
|
||||
/// [`HarnessError::SpawnFailed`] if the process cannot start. The
|
||||
/// underlying [`std::io::Error`] is deliberately reduced to its
|
||||
/// [`std::io::ErrorKind`] so a filesystem path cannot travel in an error.
|
||||
pub(crate) fn spawn(program: &str, arguments: &[&str]) -> Result<Self, HarnessError> {
|
||||
if program.trim().is_empty() {
|
||||
return Err(HarnessError::EmptyProgram);
|
||||
}
|
||||
let mut command = Command::new(program);
|
||||
command
|
||||
.args(arguments)
|
||||
.env_clear()
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::null());
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::process::CommandExt;
|
||||
// The child leads its own group, so its group ID is its process ID
|
||||
// and every descendant it does not deliberately detach joins it.
|
||||
command.process_group(0);
|
||||
}
|
||||
for name in PROBE_ENV_ALLOWLIST {
|
||||
if let Some(value) = std::env::var_os(name) {
|
||||
command.env(name, value);
|
||||
}
|
||||
}
|
||||
let mut child = command
|
||||
.spawn()
|
||||
.map_err(|error| HarnessError::SpawnFailed(error.kind()))?;
|
||||
let broken = HarnessError::SpawnFailed(std::io::ErrorKind::BrokenPipe);
|
||||
let group = child.id();
|
||||
let stdin = child.stdin.take().ok_or(broken)?;
|
||||
let stdout = child.stdout.take().ok_or(broken)?;
|
||||
Ok(Self {
|
||||
child: Arc::new(Mutex::new(child)),
|
||||
group,
|
||||
stdin,
|
||||
stdout,
|
||||
})
|
||||
}
|
||||
|
||||
/// Splits the child into its three independent owners.
|
||||
pub(crate) fn into_parts(self) -> (ProbeWriter, ProbeReader, ProbeKiller) {
|
||||
(
|
||||
ProbeWriter { stdin: self.stdin },
|
||||
ProbeReader {
|
||||
reader: BufReader::new(self.stdout),
|
||||
},
|
||||
ProbeKiller {
|
||||
child: self.child,
|
||||
group: self.group,
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// The writing half.
|
||||
pub(crate) struct ProbeWriter {
|
||||
stdin: ChildStdin,
|
||||
}
|
||||
|
||||
impl ProbeWriter {
|
||||
/// Writes one newline-terminated frame.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`HarnessError::Eof`] once the child has closed its input.
|
||||
pub(crate) fn write_line(&mut self, line: &str) -> Result<(), HarnessError> {
|
||||
self.stdin
|
||||
.write_all(line.as_bytes())
|
||||
.and_then(|()| self.stdin.write_all(b"\n"))
|
||||
.and_then(|()| self.stdin.flush())
|
||||
.map_err(|_| HarnessError::Eof)
|
||||
}
|
||||
}
|
||||
|
||||
/// The reading half. Its methods block, so it belongs on its own thread.
|
||||
pub(crate) struct ProbeReader {
|
||||
reader: BufReader<ChildStdout>,
|
||||
}
|
||||
|
||||
impl ProbeReader {
|
||||
/// Reads one frame, refusing an oversized line instead of buffering it.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`HarnessError::Eof`] at end of stream and
|
||||
/// [`HarnessError::OversizedLine`] past [`MAX_LINE_BYTES`].
|
||||
pub(crate) fn read_line(&mut self) -> Result<String, HarnessError> {
|
||||
let mut buffer = Vec::new();
|
||||
let read = self
|
||||
.reader
|
||||
.by_ref()
|
||||
.take(MAX_LINE_BYTES as u64 + 1)
|
||||
.read_until(b'\n', &mut buffer)
|
||||
.map_err(|_| HarnessError::Eof)?;
|
||||
if read == 0 {
|
||||
return Err(HarnessError::Eof);
|
||||
}
|
||||
if read > MAX_LINE_BYTES {
|
||||
return Err(HarnessError::OversizedLine(read));
|
||||
}
|
||||
String::from_utf8(buffer).map_err(|_| HarnessError::Malformed)
|
||||
}
|
||||
}
|
||||
|
||||
/// The half a supervisor keeps so it can always stop the child.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ProbeKiller {
|
||||
child: Arc<Mutex<Child>>,
|
||||
group: u32,
|
||||
}
|
||||
|
||||
impl ProbeKiller {
|
||||
/// Ends the child and everything in its process group, then reaps it.
|
||||
///
|
||||
/// Killing the whole group is what closes every write end of our stdout
|
||||
/// pipe, which is what unblocks the reader. Idempotent: killing an
|
||||
/// already-dead process is not an error worth reporting, and a poisoned
|
||||
/// lock means another thread died mid-kill, which the operating system has
|
||||
/// already resolved.
|
||||
pub(crate) fn kill(&self) {
|
||||
self.kill_group();
|
||||
if let Ok(mut child) = self.child.lock() {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn kill_group(&self) {
|
||||
use nix::sys::signal::{Signal, killpg};
|
||||
use nix::unistd::Pid;
|
||||
|
||||
if let Ok(group) = i32::try_from(self.group) {
|
||||
let _ = killpg(Pid::from_raw(group), Signal::SIGKILL);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn kill_group(&self) {}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{MAX_LINE_BYTES, PROBE_ENV_ALLOWLIST, ProbeChild};
|
||||
use crate::HarnessError;
|
||||
|
||||
#[test]
|
||||
fn an_empty_program_is_refused_before_spawning() {
|
||||
assert!(matches!(
|
||||
ProbeChild::spawn(" ", &[]).err(),
|
||||
Some(HarnessError::EmptyProgram)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_missing_program_reports_a_kind_and_never_a_path() {
|
||||
let Err(error) = ProbeChild::spawn("lumbridge-no-such-probe-binary", &[]) else {
|
||||
panic!("a missing program cannot spawn");
|
||||
};
|
||||
assert!(matches!(error, HarnessError::SpawnFailed(_)));
|
||||
assert!(
|
||||
!error.to_string().contains("lumbridge-no-such-probe-binary"),
|
||||
"an error must not carry a filesystem path"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_allowlist_is_literal_sorted_and_free_of_secret_bearing_names() {
|
||||
let mut sorted = PROBE_ENV_ALLOWLIST.to_vec();
|
||||
sorted.sort_unstable();
|
||||
assert_eq!(sorted, PROBE_ENV_ALLOWLIST);
|
||||
for name in PROBE_ENV_ALLOWLIST {
|
||||
let upper = name.to_uppercase();
|
||||
assert!(
|
||||
!upper.contains("TOKEN")
|
||||
&& !upper.contains("KEY")
|
||||
&& !upper.contains("SECRET")
|
||||
&& !upper.contains("PASSWORD"),
|
||||
"{name} could carry a credential into a child"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_child_round_trips_one_frame() {
|
||||
let Ok(child) = ProbeChild::spawn("cat", &[]) else {
|
||||
return; // cat is not guaranteed on every supported platform.
|
||||
};
|
||||
let (mut writer, mut reader, killer) = child.into_parts();
|
||||
writer
|
||||
.write_line(r#"{"ok":true}"#)
|
||||
.expect("cat accepts input");
|
||||
assert_eq!(
|
||||
reader.read_line().expect("cat echoes the line"),
|
||||
"{\"ok\":true}\n"
|
||||
);
|
||||
killer.kill();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_oversized_frame_is_refused_rather_than_buffered() {
|
||||
// Generated by the child so the test never deadlocks writing more than
|
||||
// a pipe buffer into a process that is not reading.
|
||||
let script = format!("head -c {} /dev/zero | tr '\\0' 'x'", MAX_LINE_BYTES + 16);
|
||||
let Ok(child) = ProbeChild::spawn("sh", &["-c", &script]) else {
|
||||
return;
|
||||
};
|
||||
let (_writer, mut reader, killer) = child.into_parts();
|
||||
assert!(matches!(
|
||||
reader.read_line(),
|
||||
Err(HarnessError::OversizedLine(_))
|
||||
));
|
||||
killer.kill();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn killing_a_child_unblocks_a_reader_blocked_on_a_grandchild() {
|
||||
// A launcher that backgrounds a helper and waits is the shape of the
|
||||
// npm Codex distribution: the helper inherits our stdout, so killing
|
||||
// only the direct child would leave the reader blocked forever.
|
||||
let Ok(child) = ProbeChild::spawn("sh", &["-c", "sleep 30 & wait"]) else {
|
||||
return;
|
||||
};
|
||||
let (_writer, mut reader, killer) = child.into_parts();
|
||||
let reading = std::thread::spawn(move || reader.read_line());
|
||||
killer.kill();
|
||||
let outcome = reading.join().expect("the reader thread must not panic");
|
||||
assert!(
|
||||
matches!(outcome, Err(HarnessError::Eof)),
|
||||
"killing the group must end the read rather than hang it"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn killing_a_child_unblocks_its_reader() {
|
||||
// `sleep` never writes, so the reader is blocked until the kill lands.
|
||||
let Ok(child) = ProbeChild::spawn("sleep", &["30"]) else {
|
||||
return;
|
||||
};
|
||||
let (_writer, mut reader, killer) = child.into_parts();
|
||||
let reading = std::thread::spawn(move || reader.read_line());
|
||||
killer.kill();
|
||||
let outcome = reading.join().expect("the reader thread must not panic");
|
||||
assert!(
|
||||
matches!(outcome, Err(HarnessError::Eof)),
|
||||
"a killed child must end the read rather than hang it"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,836 @@
|
||||
//! The Claude Code usage adapter, over two different surfaces.
|
||||
//!
|
||||
//! **Subscription windows** ([`statusline`]) are the real quota. Claude Code
|
||||
//! 2.1.80 and later pipe a `rate_limits` object carrying `five_hour` and
|
||||
//! `seven_day` — each with `used_percentage` and `resets_at` — to the
|
||||
//! configured `statusLine` command on every turn. Those are rate-limit headers
|
||||
//! the CLI already received on its own API responses, so reading them costs
|
||||
//! nothing and they are `ProviderReported`. Lumbridge does not call the
|
||||
//! account's usage endpoint for them: that would mean reading Claude Code's
|
||||
//! OAuth credential, which AGENTS.md forbids. The status line is pushed to a
|
||||
//! command the user installs, so no credential is ever touched.
|
||||
//!
|
||||
//! **Token consumption** ([`transcript`]) comes from the session transcripts
|
||||
//! Claude Code writes under `~/.claude/projects`. Those record the API's
|
||||
//! `usage` object per assistant turn. A running total is the harness's own
|
||||
//! record rather than a provider's statement of account, so it is
|
||||
//! `HarnessReported`, and it carries no ceiling — spend is not a quota.
|
||||
//!
|
||||
//! Both are tail followers. Each remembers a byte offset and reads only what
|
||||
//! was appended since, which is what makes the token total monotonic — the
|
||||
//! ledger rejects an observation that moves a profile backwards.
|
||||
|
||||
mod statusline;
|
||||
mod transcript;
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::mpsc::{self, Receiver, SyncSender, TryRecvError};
|
||||
use std::thread::{self, JoinHandle};
|
||||
use std::time::Duration;
|
||||
|
||||
use lumbridge_core::{AccountProfile, UsageObservation, UsageProvenance, UsageUnit};
|
||||
|
||||
use crate::HarnessError;
|
||||
use crate::claude::statusline::parse_feed_line;
|
||||
pub use crate::claude::statusline::{ClaudeWindowKind, WindowReading};
|
||||
use crate::claude::transcript::MAX_RECORD_BYTES;
|
||||
use crate::clock::MonotonicWallClock;
|
||||
use crate::probe::{ProbeHealth, ProbeOutcome, UsageProbe};
|
||||
|
||||
pub use transcript::TokenTally;
|
||||
|
||||
/// Transcripts are appended to constantly; polling faster than this buys
|
||||
/// nothing but disk churn.
|
||||
const MIN_POLL_INTERVAL: Duration = Duration::from_secs(1);
|
||||
/// How long the worker sleeps between filesystem scans.
|
||||
const WORKER_TICK: Duration = Duration::from_millis(250);
|
||||
/// A bound on how much of the tree one scan will walk.
|
||||
const MAX_DEPTH: usize = 6;
|
||||
const MAX_FILES_PER_SCAN: usize = 4_096;
|
||||
/// A bound on how much any one poll will read, so a large backlog is absorbed
|
||||
/// across several polls instead of stalling one.
|
||||
const MAX_BYTES_PER_POLL: usize = 8 * 1024 * 1024;
|
||||
const EVENT_QUEUE: usize = 64;
|
||||
|
||||
/// Where and how often to read transcripts.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ClaudeCodeProbeOptions {
|
||||
/// The Claude Code projects root. Defaults to `$CLAUDE_CONFIG_DIR/projects`
|
||||
/// or `~/.claude/projects`.
|
||||
pub projects_root: PathBuf,
|
||||
/// Where the Lumbridge status-line bridge appends rate-limit readings.
|
||||
/// Absent until the user installs the bridge, which is why the windows read
|
||||
/// as unavailable rather than as zero before then.
|
||||
pub rate_limit_feed: PathBuf,
|
||||
pub poll_interval: Duration,
|
||||
}
|
||||
|
||||
impl Default for ClaudeCodeProbeOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
projects_root: default_projects_root(),
|
||||
rate_limit_feed: default_rate_limit_feed(),
|
||||
poll_interval: Duration::from_secs(10),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Where the status-line bridge writes, under Lumbridge's own data root rather
|
||||
/// than inside the harness's configuration.
|
||||
///
|
||||
/// `LUMBRIDGE_CLAUDE_FEED` overrides it, and must be read here as well as in the
|
||||
/// bridge: the two halves have to resolve the same path or the bridge writes
|
||||
/// somewhere the probe never looks and the windows read as unavailable forever.
|
||||
#[must_use]
|
||||
pub fn default_rate_limit_feed() -> PathBuf {
|
||||
if let Some(configured) = std::env::var_os("LUMBRIDGE_CLAUDE_FEED") {
|
||||
return PathBuf::from(configured);
|
||||
}
|
||||
let root = std::env::var_os("XDG_DATA_HOME").map_or_else(
|
||||
|| {
|
||||
std::env::var_os("HOME")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_default()
|
||||
.join(".local")
|
||||
.join("share")
|
||||
},
|
||||
PathBuf::from,
|
||||
);
|
||||
root.join("lumbridge").join("claude-rate-limits.jsonl")
|
||||
}
|
||||
|
||||
/// The documented default location of Claude Code's session transcripts.
|
||||
#[must_use]
|
||||
pub fn default_projects_root() -> PathBuf {
|
||||
if let Some(configured) = std::env::var_os("CLAUDE_CONFIG_DIR") {
|
||||
return PathBuf::from(configured).join("projects");
|
||||
}
|
||||
std::env::var_os("HOME")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_default()
|
||||
.join(".claude")
|
||||
.join("projects")
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum WorkerEvent {
|
||||
Observation(UsageObservation),
|
||||
Health(ProbeHealth),
|
||||
}
|
||||
|
||||
/// A running Claude Code usage probe.
|
||||
pub struct ClaudeCodeProbe {
|
||||
profiles: Vec<AccountProfile>,
|
||||
events: Option<Receiver<WorkerEvent>>,
|
||||
stop: Arc<AtomicBool>,
|
||||
worker: Option<JoinHandle<()>>,
|
||||
health: ProbeHealth,
|
||||
tokens: TokenTally,
|
||||
}
|
||||
|
||||
impl ClaudeCodeProbe {
|
||||
/// The static, account-free profiles this probe reports under: one per
|
||||
/// documented subscription window, plus the transcript token total.
|
||||
///
|
||||
/// The windows come first because they are the quota — what a user means
|
||||
/// by "how much do I have left" — and the token total is the supporting
|
||||
/// fact about spend.
|
||||
fn profiles_for() -> Result<Vec<AccountProfile>, HarnessError> {
|
||||
let mut profiles = Vec::new();
|
||||
for kind in ClaudeWindowKind::ALL {
|
||||
profiles.push(
|
||||
AccountProfile::new(
|
||||
kind.profile_id(),
|
||||
"Claude Code",
|
||||
"Anthropic",
|
||||
kind.scope(),
|
||||
"subscription",
|
||||
)
|
||||
.map_err(|_| HarnessError::EmptyProgram)?,
|
||||
);
|
||||
}
|
||||
profiles.push(
|
||||
AccountProfile::new(
|
||||
"claude-code-transcripts",
|
||||
"Claude Code",
|
||||
"Anthropic",
|
||||
"session transcripts",
|
||||
"subscription",
|
||||
)
|
||||
.map_err(|_| HarnessError::EmptyProgram)?,
|
||||
);
|
||||
Ok(profiles)
|
||||
}
|
||||
|
||||
/// Starts the probe and its worker.
|
||||
///
|
||||
/// A missing projects directory is not an error: the probe reports zero
|
||||
/// and keeps looking, because Claude Code creates it on first use.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`HarnessError::ProbeIntervalTooShort`] below one second.
|
||||
pub fn start(options: ClaudeCodeProbeOptions) -> Result<Self, HarnessError> {
|
||||
if options.poll_interval < MIN_POLL_INTERVAL {
|
||||
return Err(HarnessError::ProbeIntervalTooShort);
|
||||
}
|
||||
let profiles = Self::profiles_for()?;
|
||||
let worker_profiles = profiles.clone();
|
||||
let (event_sender, events) = mpsc::sync_channel(EVENT_QUEUE);
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let worker_stop = Arc::clone(&stop);
|
||||
let worker = thread::Builder::new()
|
||||
.name("lumbridge-claude-probe".to_owned())
|
||||
.spawn(move || {
|
||||
run_worker(&options, &worker_profiles, &event_sender, &worker_stop);
|
||||
})
|
||||
.map_err(|error| HarnessError::SpawnFailed(error.kind()))?;
|
||||
|
||||
Ok(Self {
|
||||
profiles,
|
||||
events: Some(events),
|
||||
stop,
|
||||
worker: Some(worker),
|
||||
health: ProbeHealth::Starting,
|
||||
tokens: TokenTally::default(),
|
||||
})
|
||||
}
|
||||
|
||||
/// The token split behind the latest observation.
|
||||
///
|
||||
/// The ledger holds only the total; this keeps the breakdown available for
|
||||
/// a detail view without a second pass over the transcripts.
|
||||
#[must_use]
|
||||
pub const fn tokens(&self) -> TokenTally {
|
||||
self.tokens
|
||||
}
|
||||
}
|
||||
|
||||
impl UsageProbe for ClaudeCodeProbe {
|
||||
fn poll(&mut self) -> ProbeOutcome {
|
||||
let Some(events) = self.events.as_ref() else {
|
||||
return ProbeOutcome::idle(self.health);
|
||||
};
|
||||
let mut observations = Vec::new();
|
||||
loop {
|
||||
match events.try_recv() {
|
||||
Ok(WorkerEvent::Observation(observation)) => observations.push(observation),
|
||||
Ok(WorkerEvent::Health(health)) => self.health = health,
|
||||
Err(TryRecvError::Empty) => break,
|
||||
Err(TryRecvError::Disconnected) => {
|
||||
if !self.health.is_faulted() {
|
||||
self.health = ProbeHealth::Stopped;
|
||||
}
|
||||
self.events = None;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
ProbeOutcome::new(observations, self.health)
|
||||
}
|
||||
|
||||
fn profiles(&self) -> &[AccountProfile] {
|
||||
&self.profiles
|
||||
}
|
||||
|
||||
fn shutdown(&mut self) {
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
self.events = None;
|
||||
if let Some(worker) = self.worker.take() {
|
||||
let _ = worker.join();
|
||||
}
|
||||
if !self.health.is_faulted() {
|
||||
self.health = ProbeHealth::Stopped;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ClaudeCodeProbe {
|
||||
fn drop(&mut self) {
|
||||
self.shutdown();
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-transcript read position. The offset is what makes the total monotonic.
|
||||
#[derive(Default)]
|
||||
struct FollowState {
|
||||
offsets: BTreeMap<PathBuf, u64>,
|
||||
tokens: TokenTally,
|
||||
/// Read position in the status-line feed, and the newest reading seen.
|
||||
feed_offset: u64,
|
||||
windows: BTreeMap<&'static str, WindowReading>,
|
||||
}
|
||||
|
||||
/// Reads whatever the status-line bridge appended and keeps the newest line.
|
||||
///
|
||||
/// The feed is a state snapshot per line, not a stream of deltas, so only the
|
||||
/// last complete line matters — an older one is superseded, never summed.
|
||||
fn follow_feed(path: &Path, state: &mut FollowState, observed_at_ms: u64) {
|
||||
let Ok(metadata) = std::fs::metadata(path) else {
|
||||
return;
|
||||
};
|
||||
let length = metadata.len();
|
||||
if length < state.feed_offset {
|
||||
state.feed_offset = 0;
|
||||
}
|
||||
if length == state.feed_offset {
|
||||
return;
|
||||
}
|
||||
let take = usize::try_from(length - state.feed_offset)
|
||||
.unwrap_or(usize::MAX)
|
||||
.min(MAX_BYTES_PER_POLL);
|
||||
let Ok(mut file) = File::open(path) else {
|
||||
return;
|
||||
};
|
||||
if file.seek(SeekFrom::Start(state.feed_offset)).is_err() {
|
||||
return;
|
||||
}
|
||||
let mut buffer = vec![0_u8; take];
|
||||
let Ok(read) = file.read(&mut buffer) else {
|
||||
return;
|
||||
};
|
||||
buffer.truncate(read);
|
||||
let chunk = String::from_utf8_lossy(&buffer);
|
||||
|
||||
let mut consumed = 0;
|
||||
let mut newest = None;
|
||||
for line in chunk.split_inclusive('\n') {
|
||||
if !line.ends_with('\n') {
|
||||
break;
|
||||
}
|
||||
consumed += line.len();
|
||||
if line.len() <= MAX_RECORD_BYTES
|
||||
&& let Some(readings) = parse_feed_line(line.trim_end(), observed_at_ms)
|
||||
{
|
||||
newest = Some(readings);
|
||||
}
|
||||
}
|
||||
state.feed_offset += consumed as u64;
|
||||
if let Some(readings) = newest {
|
||||
for (kind, reading) in readings {
|
||||
state.windows.insert(kind.profile_id(), reading);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn run_worker(
|
||||
options: &ClaudeCodeProbeOptions,
|
||||
profiles: &[AccountProfile],
|
||||
events: &SyncSender<WorkerEvent>,
|
||||
stop: &AtomicBool,
|
||||
) {
|
||||
let token_profile = profiles
|
||||
.iter()
|
||||
.find(|profile| profile.id().as_str() == "claude-code-transcripts")
|
||||
.cloned();
|
||||
let mut clock = MonotonicWallClock::start();
|
||||
let mut state = FollowState::default();
|
||||
let poll_interval_ms = u64::try_from(options.poll_interval.as_millis()).unwrap_or(u64::MAX);
|
||||
let mut last_poll_ms = 0;
|
||||
let mut last_emitted = TokenTally::default();
|
||||
let mut last_windows: BTreeMap<&'static str, WindowReading> = BTreeMap::new();
|
||||
let mut primed = false;
|
||||
|
||||
loop {
|
||||
if stop.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
let now_ms = clock.now_ms();
|
||||
// While priming, scan every tick to work through the backlog quickly.
|
||||
// After that, respect the poll interval.
|
||||
if primed && now_ms.saturating_sub(last_poll_ms) < poll_interval_ms {
|
||||
thread::sleep(WORKER_TICK);
|
||||
continue;
|
||||
}
|
||||
last_poll_ms = now_ms;
|
||||
|
||||
let caught_up = scan(&options.projects_root, &mut state);
|
||||
// The quota windows are independent of the transcript backlog: a
|
||||
// status-line reading is a complete snapshot, so it can be reported
|
||||
// immediately rather than waiting for priming.
|
||||
follow_feed(&options.rate_limit_feed, &mut state, now_ms);
|
||||
for kind in ClaudeWindowKind::ALL {
|
||||
let Some(reading) = state.windows.get(kind.profile_id()).copied() else {
|
||||
continue;
|
||||
};
|
||||
if last_windows.get(kind.profile_id()) == Some(&reading) {
|
||||
continue;
|
||||
}
|
||||
last_windows.insert(kind.profile_id(), reading);
|
||||
let Some(profile) = profiles
|
||||
.iter()
|
||||
.find(|profile| profile.id().as_str() == kind.profile_id())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let observation = window_observation(profile, reading, now_ms);
|
||||
if matches!(
|
||||
events.try_send(WorkerEvent::Observation(observation)),
|
||||
Err(mpsc::TrySendError::Disconnected(_))
|
||||
) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Nothing is reported until the follower has read every existing
|
||||
// transcript to its end. A partially-read backlog looks exactly like
|
||||
// an enormous burst of spend, and the ledger would derive a burn rate
|
||||
// from it — one early run reported forty-six billion tokens an hour,
|
||||
// which was catch-up, not usage. The first reading must be a true
|
||||
// baseline before any rate can mean anything.
|
||||
if !caught_up {
|
||||
thread::sleep(WORKER_TICK);
|
||||
continue;
|
||||
}
|
||||
if !primed {
|
||||
primed = true;
|
||||
let _ = events.try_send(WorkerEvent::Health(ProbeHealth::Ready));
|
||||
}
|
||||
|
||||
// Only speak when the number changed. An unchanged total re-recorded
|
||||
// every poll would flush the ledger's bounded history of real readings.
|
||||
if state.tokens == last_emitted {
|
||||
continue;
|
||||
}
|
||||
last_emitted = state.tokens;
|
||||
let Some(profile) = token_profile.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
let Ok(observation) = UsageObservation::counted(
|
||||
profile.id().clone(),
|
||||
UsageUnit::Tokens,
|
||||
state.tokens.total(),
|
||||
UsageProvenance::HarnessReported,
|
||||
now_ms,
|
||||
) else {
|
||||
continue;
|
||||
};
|
||||
// No `with_limit` and no `with_window`: the transcript states what was
|
||||
// spent and says nothing about a ceiling or a reset. Attaching either
|
||||
// would be an invention.
|
||||
if matches!(
|
||||
events.try_send(WorkerEvent::Observation(observation)),
|
||||
Err(mpsc::TrySendError::Disconnected(_))
|
||||
) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Turns a window reading into the observation the ledger should hold.
|
||||
///
|
||||
/// A usable reading is `ProviderReported`: the CLI is relaying rate-limit
|
||||
/// headers from its own API responses, not computing a number. An absent
|
||||
/// window is an explicit gap — an account with no plan limits, or a window the
|
||||
/// API has stopped reporting — never a zero.
|
||||
fn window_observation(
|
||||
profile: &AccountProfile,
|
||||
reading: WindowReading,
|
||||
observed_at_ms: u64,
|
||||
) -> UsageObservation {
|
||||
let WindowReading::Usable { permille, window } = reading else {
|
||||
return UsageObservation::unavailable(profile.id().clone(), observed_at_ms);
|
||||
};
|
||||
let Ok(observation) = UsageObservation::counted(
|
||||
profile.id().clone(),
|
||||
UsageUnit::WindowPermille,
|
||||
permille,
|
||||
UsageProvenance::ProviderReported,
|
||||
observed_at_ms,
|
||||
) else {
|
||||
return UsageObservation::unavailable(profile.id().clone(), observed_at_ms);
|
||||
};
|
||||
let observation = observation
|
||||
.with_limit(1_000)
|
||||
.unwrap_or_else(|_| UsageObservation::unavailable(profile.id().clone(), observed_at_ms));
|
||||
let Some(window) = window else {
|
||||
return observation;
|
||||
};
|
||||
observation
|
||||
.clone()
|
||||
.with_window(window)
|
||||
.unwrap_or(observation)
|
||||
}
|
||||
|
||||
/// Walks the projects tree and reads whatever was appended since last time.
|
||||
///
|
||||
/// Returns whether every transcript was read to its end. A `false` means the
|
||||
/// per-poll byte budget ran out with work remaining, and the running total is
|
||||
/// therefore mid-catch-up rather than current.
|
||||
fn scan(root: &Path, state: &mut FollowState) -> bool {
|
||||
let mut transcripts = Vec::new();
|
||||
collect_transcripts(root, 0, &mut transcripts);
|
||||
let mut budget = MAX_BYTES_PER_POLL;
|
||||
let mut caught_up = true;
|
||||
for path in transcripts {
|
||||
if budget == 0 {
|
||||
return false;
|
||||
}
|
||||
if !follow(&path, state, &mut budget) {
|
||||
caught_up = false;
|
||||
}
|
||||
}
|
||||
caught_up
|
||||
}
|
||||
|
||||
fn collect_transcripts(directory: &Path, depth: usize, found: &mut Vec<PathBuf>) {
|
||||
if depth > MAX_DEPTH || found.len() >= MAX_FILES_PER_SCAN {
|
||||
return;
|
||||
}
|
||||
let Ok(entries) = std::fs::read_dir(directory) else {
|
||||
return;
|
||||
};
|
||||
for entry in entries.flatten() {
|
||||
if found.len() >= MAX_FILES_PER_SCAN {
|
||||
return;
|
||||
}
|
||||
let path = entry.path();
|
||||
let Ok(file_type) = entry.file_type() else {
|
||||
continue;
|
||||
};
|
||||
// Deliberately not following symlinks: a link inside the projects tree
|
||||
// could otherwise point this reader at an arbitrary file.
|
||||
if file_type.is_dir() {
|
||||
collect_transcripts(&path, depth + 1, found);
|
||||
} else if file_type.is_file() && path.extension().is_some_and(|ext| ext == "jsonl") {
|
||||
found.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads one transcript from its remembered offset.
|
||||
///
|
||||
/// Returns whether this file is now read to its end.
|
||||
fn follow(path: &Path, state: &mut FollowState, budget: &mut usize) -> bool {
|
||||
let offset = state.offsets.get(path).copied().unwrap_or(0);
|
||||
let Ok(metadata) = std::fs::metadata(path) else {
|
||||
return true;
|
||||
};
|
||||
let length = metadata.len();
|
||||
if length < offset {
|
||||
// The file shrank, so it is not the file we were reading. Start over
|
||||
// rather than trusting an offset into different content.
|
||||
state.offsets.insert(path.to_path_buf(), 0);
|
||||
return false;
|
||||
}
|
||||
let pending = usize::try_from(length - offset).unwrap_or(usize::MAX);
|
||||
if pending == 0 {
|
||||
return true;
|
||||
}
|
||||
let take = pending.min(*budget);
|
||||
|
||||
let Ok(mut file) = File::open(path) else {
|
||||
return true;
|
||||
};
|
||||
if file.seek(SeekFrom::Start(offset)).is_err() {
|
||||
return true;
|
||||
}
|
||||
let mut buffer = vec![0_u8; take];
|
||||
let Ok(read) = file.read(&mut buffer) else {
|
||||
return true;
|
||||
};
|
||||
buffer.truncate(read);
|
||||
// A transcript is UTF-8 JSON. A chunk boundary can split a multi-byte
|
||||
// character, so decode lossily and stop at the last complete line — the
|
||||
// offset only advances past whole records either way.
|
||||
let chunk = String::from_utf8_lossy(&buffer);
|
||||
let (tally, consumed) = transcript::consume_chunk(&chunk);
|
||||
state.tokens.add(tally.tokens);
|
||||
state
|
||||
.offsets
|
||||
.insert(path.to_path_buf(), offset + consumed as u64);
|
||||
*budget = budget.saturating_sub(consumed);
|
||||
// Caught up only if we took the whole remainder and parsed all of it. A
|
||||
// trailing partial line counts as caught up: the record is still being
|
||||
// written, and waiting for it is correct.
|
||||
take == pending && pending.saturating_sub(consumed) < MAX_RECORD_BYTES
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
ClaudeCodeProbe, ClaudeCodeProbeOptions, FollowState, TokenTally, default_projects_root,
|
||||
scan,
|
||||
};
|
||||
use crate::HarnessError;
|
||||
use crate::probe::{ProbeHealth, UsageProbe};
|
||||
use lumbridge_core::{UsageProvenance, UsageUnit};
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
|
||||
const RECORD: &str = r#"{"type":"assistant","message":{"model":"claude-opus-5","content":[{"type":"text","text":"private"}],"usage":{"input_tokens":1,"output_tokens":10,"cache_creation_input_tokens":100,"cache_read_input_tokens":1000}}}"#;
|
||||
|
||||
fn scratch(tag: &str) -> PathBuf {
|
||||
let root = std::env::temp_dir().join(format!(
|
||||
"lumbridge-claude-probe-{}-{tag}",
|
||||
std::process::id()
|
||||
));
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
std::fs::create_dir_all(root.join("project-a")).expect("scratch dir");
|
||||
root
|
||||
}
|
||||
|
||||
fn append(path: &PathBuf, times: usize) {
|
||||
let mut file = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(path)
|
||||
.expect("transcript is writable");
|
||||
for _ in 0..times {
|
||||
writeln!(file, "{RECORD}").expect("append");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_missing_projects_root_reports_zero_rather_than_failing() {
|
||||
let mut state = FollowState::default();
|
||||
scan(
|
||||
&PathBuf::from("/nonexistent/lumbridge/claude/root"),
|
||||
&mut state,
|
||||
);
|
||||
assert_eq!(state.tokens, TokenTally::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn appended_records_accumulate_and_never_double_count() {
|
||||
let root = scratch("accumulate");
|
||||
let transcript = root.join("project-a").join("session.jsonl");
|
||||
append(&transcript, 2);
|
||||
|
||||
let mut state = FollowState::default();
|
||||
assert!(scan(&root, &mut state));
|
||||
assert_eq!(state.tokens.total(), 2 * 1_111);
|
||||
|
||||
// A second scan with no new bytes must add nothing.
|
||||
assert!(scan(&root, &mut state));
|
||||
assert_eq!(state.tokens.total(), 2 * 1_111);
|
||||
|
||||
// Only the newly appended record is counted.
|
||||
append(&transcript, 1);
|
||||
assert!(scan(&root, &mut state));
|
||||
assert_eq!(state.tokens.total(), 3 * 1_111);
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nested_subagent_transcripts_are_counted_too() {
|
||||
let root = scratch("nested");
|
||||
let nested = root.join("project-a").join("session").join("subagents");
|
||||
std::fs::create_dir_all(&nested).expect("nested dirs");
|
||||
append(&root.join("project-a").join("session.jsonl"), 1);
|
||||
append(&nested.join("agent-1.jsonl"), 1);
|
||||
|
||||
let mut state = FollowState::default();
|
||||
assert!(scan(&root, &mut state));
|
||||
assert_eq!(
|
||||
state.tokens.total(),
|
||||
2 * 1_111,
|
||||
"a subagent's tokens are the user's tokens"
|
||||
);
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_truncated_transcript_restarts_rather_than_reading_a_stale_offset() {
|
||||
let root = scratch("truncated");
|
||||
let transcript = root.join("project-a").join("session.jsonl");
|
||||
append(&transcript, 3);
|
||||
let mut state = FollowState::default();
|
||||
assert!(scan(&root, &mut state));
|
||||
let before = state.tokens.total();
|
||||
assert_eq!(before, 3 * 1_111);
|
||||
|
||||
std::fs::write(&transcript, "").expect("truncate");
|
||||
// The shrink is reported as not-caught-up on purpose: the offset was
|
||||
// just reset, so the probe must re-scan before it trusts the total
|
||||
// again rather than reporting mid-reset.
|
||||
assert!(
|
||||
!scan(&root, &mut state),
|
||||
"a shrunk file asks to be read again"
|
||||
);
|
||||
assert!(scan(&root, &mut state), "and is caught up on the next pass");
|
||||
assert_eq!(
|
||||
state.tokens.total(),
|
||||
before,
|
||||
"a shrunk file resets its offset without inventing or losing tokens"
|
||||
);
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_transcript_files_are_ignored() {
|
||||
let root = scratch("ignored");
|
||||
std::fs::write(root.join("project-a").join("notes.md"), RECORD).expect("write");
|
||||
std::fs::write(root.join("project-a").join("history.json"), RECORD).expect("write");
|
||||
let mut state = FollowState::default();
|
||||
assert!(scan(&root, &mut state));
|
||||
assert_eq!(state.tokens, TokenTally::default());
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_too_frequent_poll_interval_is_refused() {
|
||||
let options = ClaudeCodeProbeOptions {
|
||||
poll_interval: Duration::from_millis(10),
|
||||
..ClaudeCodeProbeOptions::default()
|
||||
};
|
||||
assert_eq!(
|
||||
ClaudeCodeProbe::start(options).err(),
|
||||
Some(HarnessError::ProbeIntervalTooShort)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_probe_reports_harness_reported_tokens_with_no_ceiling() {
|
||||
let root = scratch("probe");
|
||||
append(&root.join("project-a").join("session.jsonl"), 4);
|
||||
let mut probe = ClaudeCodeProbe::start(ClaudeCodeProbeOptions {
|
||||
projects_root: root.clone(),
|
||||
rate_limit_feed: root.join("feed.jsonl"),
|
||||
poll_interval: Duration::from_secs(1),
|
||||
})
|
||||
.expect("the probe starts");
|
||||
|
||||
let mut observations = Vec::new();
|
||||
for _ in 0..200 {
|
||||
observations.extend(probe.poll().into_observations());
|
||||
if !observations.is_empty() {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
probe.shutdown();
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
|
||||
let observation = observations.first().expect("a reading arrives");
|
||||
assert_eq!(observation.provenance(), UsageProvenance::HarnessReported);
|
||||
assert_eq!(observation.unit(), Some(UsageUnit::Tokens));
|
||||
assert_eq!(observation.consumed(), 4 * 1_111);
|
||||
assert_eq!(
|
||||
observation.limit(),
|
||||
None,
|
||||
"a transcript states spend, never a ceiling"
|
||||
);
|
||||
assert_eq!(
|
||||
observation.window(),
|
||||
None,
|
||||
"a transcript states spend, never a reset"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_status_line_feed_becomes_provider_reported_windows() {
|
||||
let root = scratch("windows");
|
||||
let feed = root.join("feed.jsonl");
|
||||
// The shape the bridge writes, with a reset far enough ahead that the
|
||||
// window is still open when the probe reads it.
|
||||
let resets_at = (std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("a clock after the epoch")
|
||||
.as_secs())
|
||||
+ 3_600;
|
||||
std::fs::write(
|
||||
&feed,
|
||||
format!(
|
||||
"{{\"rate_limits_available\":true,\"rate_limits\":{{\"five_hour\":{{\"used_percentage\":42,\"resets_at\":{resets_at}}},\"seven_day\":{{\"used_percentage\":8.5,\"resets_at\":{resets_at}}}}}}}\n"
|
||||
),
|
||||
)
|
||||
.expect("feed is writable");
|
||||
|
||||
let mut probe = ClaudeCodeProbe::start(ClaudeCodeProbeOptions {
|
||||
projects_root: root.clone(),
|
||||
rate_limit_feed: feed,
|
||||
poll_interval: Duration::from_secs(1),
|
||||
})
|
||||
.expect("the probe starts");
|
||||
|
||||
let mut observations = Vec::new();
|
||||
for _ in 0..200 {
|
||||
observations.extend(probe.poll().into_observations());
|
||||
if observations.len() >= 2 {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
probe.shutdown();
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
|
||||
let five = observations
|
||||
.iter()
|
||||
.find(|observation| observation.profile().as_str() == "claude-code-five-hour")
|
||||
.expect("the five-hour window is reported");
|
||||
assert_eq!(
|
||||
five.provenance(),
|
||||
UsageProvenance::ProviderReported,
|
||||
"the CLI relays the provider's own rate-limit headers"
|
||||
);
|
||||
assert_eq!(five.unit(), Some(UsageUnit::WindowPermille));
|
||||
assert_eq!(five.consumed(), 420);
|
||||
assert_eq!(
|
||||
five.limit(),
|
||||
Some(1_000),
|
||||
"a percentage is a share of one whole"
|
||||
);
|
||||
assert!(five.window().is_some(), "a documented window has a reset");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_account_without_plan_limits_reports_a_gap_not_a_zero() {
|
||||
let root = scratch("nolimits");
|
||||
let feed = root.join("feed.jsonl");
|
||||
std::fs::write(&feed, "{\"rate_limits_available\":false}\n").expect("writable");
|
||||
let mut probe = ClaudeCodeProbe::start(ClaudeCodeProbeOptions {
|
||||
projects_root: root.clone(),
|
||||
rate_limit_feed: feed,
|
||||
poll_interval: Duration::from_secs(1),
|
||||
})
|
||||
.expect("starts");
|
||||
let mut observations = Vec::new();
|
||||
for _ in 0..200 {
|
||||
observations.extend(probe.poll().into_observations());
|
||||
if !observations.is_empty() {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
probe.shutdown();
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
let window = observations
|
||||
.iter()
|
||||
.find(|observation| observation.profile().as_str() == "claude-code-five-hour")
|
||||
.expect("the window is still named");
|
||||
assert_eq!(window.provenance(), UsageProvenance::Unavailable);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shutdown_is_prompt_and_idempotent() {
|
||||
let root = scratch("shutdown");
|
||||
let mut probe = ClaudeCodeProbe::start(ClaudeCodeProbeOptions {
|
||||
projects_root: root.clone(),
|
||||
rate_limit_feed: root.join("feed.jsonl"),
|
||||
poll_interval: Duration::from_secs(60),
|
||||
})
|
||||
.expect("starts");
|
||||
let started = std::time::Instant::now();
|
||||
probe.shutdown();
|
||||
probe.shutdown();
|
||||
assert!(started.elapsed() < Duration::from_secs(5));
|
||||
assert!(matches!(
|
||||
probe.poll().health(),
|
||||
ProbeHealth::Stopped | ProbeHealth::Starting | ProbeHealth::Ready
|
||||
));
|
||||
let _ = std::fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_default_root_follows_the_documented_location() {
|
||||
let root = default_projects_root();
|
||||
assert!(root.ends_with("projects"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
//! Claude Code's status-line rate-limit payload: the subscription window.
|
||||
//!
|
||||
//! Claude Code 2.1.80 and later pipe a `rate_limits` object to the configured
|
||||
//! `statusLine` command on every turn. The shape is documented inside the CLI
|
||||
//! itself, which describes `five_hour` and `seven_day` as
|
||||
//!
|
||||
//! ```text
|
||||
//! "five_hour": { // present only while the API reports it and its resets_at has not passed
|
||||
//! "used_percentage": number, // Percentage of limit used (0-100)
|
||||
//! "resets_at": number // Unix epoch seconds when this window resets
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! These are the real quota windows, and they cost nothing to read: the CLI is
|
||||
//! relaying rate-limit headers that already came back on its own API responses.
|
||||
//! That makes them `ProviderReported` under decision 0013's test — the harness
|
||||
//! forwards the provider's number rather than computing one.
|
||||
//!
|
||||
//! `rate_limits_available` is the honest-gap signal, and the CLI documents it:
|
||||
//! "False when plan rate limits do not apply (API key, Bedrock, Vertex, or
|
||||
//! missing profile scope) — `rate_limits` will be null." An API-key user has no
|
||||
//! subscription window, and this field says so rather than leaving us guessing.
|
||||
//!
|
||||
//! Nothing else from the status-line payload is modelled. It also carries the
|
||||
//! session's cost, its transcript path, the current working directory, and the
|
||||
//! model — none of which this parser can represent.
|
||||
|
||||
use lumbridge_core::UsageWindow;
|
||||
use serde::Deserialize;
|
||||
|
||||
/// Percent is 0-100 on the wire; a spend limit may exceed 100 once breached.
|
||||
const MAX_PERCENT: f64 = 100.0;
|
||||
const PERMILLE_PER_PERCENT: f64 = 10.0;
|
||||
const MILLIS_PER_SECOND: i64 = 1_000;
|
||||
/// Documented window lengths, used to give each reading a window start so the
|
||||
/// ledger can scope a burn rate to one quota period.
|
||||
const FIVE_HOUR_MS: u64 = 5 * 60 * 60 * 1_000;
|
||||
const SEVEN_DAY_MS: u64 = 7 * 24 * 60 * 60 * 1_000;
|
||||
|
||||
/// One line of the feed a Lumbridge status-line bridge appends.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct FeedLineWire {
|
||||
#[serde(default)]
|
||||
pub(crate) rate_limits_available: Option<bool>,
|
||||
#[serde(default)]
|
||||
pub(crate) rate_limits: Option<RateLimitsWire>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
pub(crate) struct RateLimitsWire {
|
||||
#[serde(default)]
|
||||
pub(crate) five_hour: Option<WindowWire>,
|
||||
#[serde(default)]
|
||||
pub(crate) seven_day: Option<WindowWire>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct WindowWire {
|
||||
/// The documented field. `utilization` is accepted as a sibling because the
|
||||
/// OAuth usage endpoint spells the same quantity that way, and a rename
|
||||
/// should degrade rather than silently go dark.
|
||||
#[serde(default)]
|
||||
pub(crate) used_percentage: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub(crate) utilization: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub(crate) resets_at: Option<i64>,
|
||||
}
|
||||
|
||||
/// Which documented subscription window a reading describes.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum ClaudeWindowKind {
|
||||
FiveHour,
|
||||
SevenDay,
|
||||
}
|
||||
|
||||
impl ClaudeWindowKind {
|
||||
pub(crate) const ALL: [Self; 2] = [Self::FiveHour, Self::SevenDay];
|
||||
|
||||
/// Static, account-free profile identifiers.
|
||||
pub(crate) const fn profile_id(self) -> &'static str {
|
||||
match self {
|
||||
Self::FiveHour => "claude-code-five-hour",
|
||||
Self::SevenDay => "claude-code-seven-day",
|
||||
}
|
||||
}
|
||||
|
||||
/// Named for the window, not for a model: these limits are account-scoped.
|
||||
pub(crate) const fn scope(self) -> &'static str {
|
||||
match self {
|
||||
Self::FiveHour => "five-hour window",
|
||||
Self::SevenDay => "seven-day window",
|
||||
}
|
||||
}
|
||||
|
||||
const fn length_ms(self) -> u64 {
|
||||
match self {
|
||||
Self::FiveHour => FIVE_HOUR_MS,
|
||||
Self::SevenDay => SEVEN_DAY_MS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What one window in the feed means at a point in time.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum WindowReading {
|
||||
/// A usable share of the window.
|
||||
Usable {
|
||||
permille: u64,
|
||||
window: Option<UsageWindow>,
|
||||
},
|
||||
/// The window is absent, or the account has no plan limits at all.
|
||||
Absent,
|
||||
}
|
||||
|
||||
/// Reads one window against the moment it was observed.
|
||||
///
|
||||
/// Total: every input produces a decision and none of them panics.
|
||||
pub(crate) fn parse_window(
|
||||
wire: &WindowWire,
|
||||
kind: ClaudeWindowKind,
|
||||
observed_at_ms: u64,
|
||||
) -> WindowReading {
|
||||
let Some(percent) = wire.used_percentage.or(wire.utilization) else {
|
||||
return WindowReading::Absent;
|
||||
};
|
||||
if !percent.is_finite() || percent < 0.0 {
|
||||
return WindowReading::Absent;
|
||||
}
|
||||
// A spend limit can report over 100 once breached; a quota share cannot
|
||||
// exceed the whole, so clamp rather than letting the ledger hold >100%.
|
||||
//
|
||||
// The value is clamped to 0..=100 and scaled by ten before rounding, so it
|
||||
// is a whole number in 0..=1000 — well inside u64 and inside the range f64
|
||||
// represents exactly. The conversion cannot truncate or lose a sign here.
|
||||
#[expect(
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss,
|
||||
reason = "a rounded whole number clamped to 0..=1000"
|
||||
)]
|
||||
let permille = (percent.clamp(0.0, MAX_PERCENT) * PERMILLE_PER_PERCENT).round() as u64;
|
||||
|
||||
let window = wire
|
||||
.resets_at
|
||||
.and_then(unix_seconds_to_millis)
|
||||
// The CLI documents that a window is present only while its resets_at
|
||||
// has not passed. One that has is stale, so its start-and-reset pair is
|
||||
// dropped and only the percentage survives.
|
||||
.filter(|resets_at_ms| *resets_at_ms > observed_at_ms)
|
||||
.and_then(|resets_at_ms| {
|
||||
let started_at_ms = resets_at_ms.saturating_sub(kind.length_ms());
|
||||
if started_at_ms <= observed_at_ms {
|
||||
UsageWindow::new(started_at_ms, resets_at_ms).ok()
|
||||
} else {
|
||||
// The local clock disagrees with the provider about the start.
|
||||
// Keep the reset, which is the half that bounds a forecast.
|
||||
Some(UsageWindow::until(resets_at_ms))
|
||||
}
|
||||
});
|
||||
|
||||
WindowReading::Usable { permille, window }
|
||||
}
|
||||
|
||||
fn unix_seconds_to_millis(seconds: i64) -> Option<u64> {
|
||||
u64::try_from(seconds.checked_mul(MILLIS_PER_SECOND)?).ok()
|
||||
}
|
||||
|
||||
/// Reads one feed line into a reading per documented window.
|
||||
///
|
||||
/// `rate_limits_available: false` is a positive statement that this account has
|
||||
/// no plan limits, so every window reads as absent rather than unknown.
|
||||
pub(crate) fn parse_feed_line(
|
||||
line: &str,
|
||||
observed_at_ms: u64,
|
||||
) -> Option<[(ClaudeWindowKind, WindowReading); 2]> {
|
||||
let feed: FeedLineWire = serde_json::from_str(line).ok()?;
|
||||
if feed.rate_limits_available == Some(false) {
|
||||
return Some([
|
||||
(ClaudeWindowKind::FiveHour, WindowReading::Absent),
|
||||
(ClaudeWindowKind::SevenDay, WindowReading::Absent),
|
||||
]);
|
||||
}
|
||||
let limits = feed.rate_limits?;
|
||||
let five = limits
|
||||
.five_hour
|
||||
.as_ref()
|
||||
.map_or(WindowReading::Absent, |w| {
|
||||
parse_window(w, ClaudeWindowKind::FiveHour, observed_at_ms)
|
||||
});
|
||||
let seven = limits
|
||||
.seven_day
|
||||
.as_ref()
|
||||
.map_or(WindowReading::Absent, |w| {
|
||||
parse_window(w, ClaudeWindowKind::SevenDay, observed_at_ms)
|
||||
});
|
||||
Some([
|
||||
(ClaudeWindowKind::FiveHour, five),
|
||||
(ClaudeWindowKind::SevenDay, seven),
|
||||
])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{ClaudeWindowKind, WindowReading, parse_feed_line};
|
||||
|
||||
const HOUR_MS: u64 = 3_600_000;
|
||||
const RESETS_AT_SECONDS: i64 = 1_800_000_000;
|
||||
const RESETS_AT_MS: u64 = 1_800_000_000_000;
|
||||
|
||||
fn line(five: &str, seven: &str) -> String {
|
||||
format!(
|
||||
r#"{{"rate_limits_available":true,"rate_limits":{{"five_hour":{five},"seven_day":{seven}}}}}"#
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_documented_shape_yields_both_windows() {
|
||||
let body = line(
|
||||
&format!(r#"{{"used_percentage":42,"resets_at":{RESETS_AT_SECONDS}}}"#),
|
||||
&format!(r#"{{"used_percentage":8.5,"resets_at":{RESETS_AT_SECONDS}}}"#),
|
||||
);
|
||||
let readings = parse_feed_line(&body, RESETS_AT_MS - HOUR_MS).expect("a valid feed line");
|
||||
let WindowReading::Usable { permille, window } = readings[0].1 else {
|
||||
panic!("the five-hour window is usable");
|
||||
};
|
||||
assert_eq!(readings[0].0, ClaudeWindowKind::FiveHour);
|
||||
assert_eq!(permille, 420);
|
||||
let window = window.expect("a documented window length builds a window");
|
||||
assert_eq!(window.resets_at_ms(), RESETS_AT_MS);
|
||||
assert_eq!(
|
||||
window.started_at_ms(),
|
||||
Some(RESETS_AT_MS - 5 * HOUR_MS),
|
||||
"the five-hour window starts five hours before it resets"
|
||||
);
|
||||
|
||||
let WindowReading::Usable { permille, .. } = readings[1].1 else {
|
||||
panic!("the seven-day window is usable");
|
||||
};
|
||||
assert_eq!(permille, 85, "a fractional percent keeps its tenth");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_api_key_account_reports_absent_rather_than_zero() {
|
||||
// The CLI documents rate_limits_available as false for API key,
|
||||
// Bedrock, Vertex, or a missing profile scope.
|
||||
let readings = parse_feed_line(
|
||||
r#"{"rate_limits_available":false,"rate_limits":null}"#,
|
||||
RESETS_AT_MS,
|
||||
)
|
||||
.expect("a valid feed line");
|
||||
assert!(readings.iter().all(|(_, r)| *r == WindowReading::Absent));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_missing_window_is_absent_without_affecting_the_other() {
|
||||
let body = line(
|
||||
&format!(r#"{{"used_percentage":10,"resets_at":{RESETS_AT_SECONDS}}}"#),
|
||||
"null",
|
||||
);
|
||||
let readings = parse_feed_line(&body, RESETS_AT_MS - HOUR_MS).expect("valid");
|
||||
assert!(matches!(readings[0].1, WindowReading::Usable { .. }));
|
||||
assert_eq!(readings[1].1, WindowReading::Absent);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_reset_that_has_passed_keeps_the_percentage_and_drops_the_window() {
|
||||
let body = line(
|
||||
&format!(r#"{{"used_percentage":99,"resets_at":{RESETS_AT_SECONDS}}}"#),
|
||||
"null",
|
||||
);
|
||||
let readings = parse_feed_line(&body, RESETS_AT_MS + 1).expect("valid");
|
||||
assert_eq!(
|
||||
readings[0].1,
|
||||
WindowReading::Usable {
|
||||
permille: 990,
|
||||
window: None
|
||||
},
|
||||
"a window whose reset has passed is stale, not current"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_oauth_spelling_of_the_same_field_is_accepted() {
|
||||
let body = line(
|
||||
&format!(r#"{{"utilization":25,"resets_at":{RESETS_AT_SECONDS}}}"#),
|
||||
"null",
|
||||
);
|
||||
let readings = parse_feed_line(&body, RESETS_AT_MS - HOUR_MS).expect("valid");
|
||||
assert!(matches!(
|
||||
readings[0].1,
|
||||
WindowReading::Usable { permille: 250, .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_over_range_or_nonsense_percentage_is_clamped_or_refused() {
|
||||
let over = line(r#"{"used_percentage":140}"#, "null");
|
||||
assert!(matches!(
|
||||
parse_feed_line(&over, 0).expect("valid")[0].1,
|
||||
WindowReading::Usable {
|
||||
permille: 1_000,
|
||||
window: None
|
||||
}
|
||||
));
|
||||
for bad in ["-1", "null"] {
|
||||
let body = line(&format!(r#"{{"used_percentage":{bad}}}"#), "null");
|
||||
assert_eq!(
|
||||
parse_feed_line(&body, 0).expect("valid")[0].1,
|
||||
WindowReading::Absent
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_or_unrelated_lines_are_refused() {
|
||||
for body in ["", "not json", "{}", r#"{"rate_limits":{}}"#] {
|
||||
let parsed = parse_feed_line(body, 0);
|
||||
assert!(
|
||||
parsed.is_none_or(|readings| readings
|
||||
.iter()
|
||||
.all(|(_, r)| *r == WindowReading::Absent)),
|
||||
"{body:?} must not produce a number"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn window_kinds_have_distinct_account_free_identifiers() {
|
||||
assert_ne!(
|
||||
ClaudeWindowKind::FiveHour.profile_id(),
|
||||
ClaudeWindowKind::SevenDay.profile_id()
|
||||
);
|
||||
for kind in ClaudeWindowKind::ALL {
|
||||
assert!(!kind.profile_id().contains('@'));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
//! The Claude Code transcript reader: token counts in, nothing else out.
|
||||
//!
|
||||
//! Claude Code writes one JSONL record per event to
|
||||
//! `~/.claude/projects/<encoded-cwd>/<session>.jsonl`. Assistant records carry
|
||||
//! the API's `usage` object verbatim, which is the only documented local
|
||||
//! surface that reports what the harness actually spent.
|
||||
//!
|
||||
//! **These files contain the user's conversations.** This module is built so
|
||||
//! that reading one cannot surface their content: the wire types below model
|
||||
//! the record type, the model name, and the four token counters, and nothing
|
||||
//! else. Serde discards every unmodelled field while parsing, so message text,
|
||||
//! tool inputs, and file contents are never deserialized into a Lumbridge value
|
||||
//! at all. That is a structural guarantee, not a convention — there is no
|
||||
//! field here that could hold them.
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
/// A transcript line is a single JSON object. Anything longer than this is not
|
||||
/// a record we can use, and buffering it would let one file grow this process.
|
||||
pub(crate) const MAX_RECORD_BYTES: usize = 1024 * 1024;
|
||||
|
||||
/// The only fields this crate reads from a transcript record.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct RecordWire {
|
||||
#[serde(rename = "type")]
|
||||
pub(crate) record_type: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) message: Option<MessageWire>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct MessageWire {
|
||||
#[serde(default)]
|
||||
pub(crate) model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(crate) usage: Option<UsageWire>,
|
||||
}
|
||||
|
||||
/// The API usage object as Claude Code records it.
|
||||
///
|
||||
/// `cache_creation`, `server_tool_use`, `iterations`, `service_tier`, `speed`,
|
||||
/// and `inference_geo` are all present on the wire and all deliberately
|
||||
/// unmodelled: a probe should not be able to read account or routing metadata
|
||||
/// it has no use for.
|
||||
///
|
||||
/// The wire names are pinned with explicit renames so the Rust field names can
|
||||
/// read naturally without the contract drifting from what Claude Code writes.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
pub(crate) struct UsageWire {
|
||||
#[serde(default, rename = "input_tokens")]
|
||||
pub(crate) input: u64,
|
||||
#[serde(default, rename = "output_tokens")]
|
||||
pub(crate) output: u64,
|
||||
#[serde(default, rename = "cache_creation_input_tokens")]
|
||||
pub(crate) cache_creation: u64,
|
||||
#[serde(default, rename = "cache_read_input_tokens")]
|
||||
pub(crate) cache_read: u64,
|
||||
}
|
||||
|
||||
/// Tokens the harness recorded, kept split by kind.
|
||||
///
|
||||
/// The four counters are priced very differently and mean different things, so
|
||||
/// they are carried separately rather than collapsed at the point of parsing.
|
||||
/// The ledger holds one number per profile; [`TokenTally::total`] is what it
|
||||
/// gets, and the split stays available for a detail view.
|
||||
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
|
||||
pub struct TokenTally {
|
||||
pub input: u64,
|
||||
pub output: u64,
|
||||
pub cache_creation: u64,
|
||||
pub cache_read: u64,
|
||||
}
|
||||
|
||||
impl TokenTally {
|
||||
/// Every token the API processed.
|
||||
///
|
||||
/// Cache reads are included because they are billed and do count against a
|
||||
/// subscription window, even though they are re-reads of context already
|
||||
/// counted once. Excluding them would understate consumption; weighting
|
||||
/// them by price would require a price list this adapter does not have and
|
||||
/// would turn a counted fact into an estimate.
|
||||
#[must_use]
|
||||
pub const fn total(self) -> u64 {
|
||||
self.input
|
||||
.saturating_add(self.output)
|
||||
.saturating_add(self.cache_creation)
|
||||
.saturating_add(self.cache_read)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn is_empty(self) -> bool {
|
||||
self.total() == 0
|
||||
}
|
||||
|
||||
pub(crate) fn add(&mut self, other: Self) {
|
||||
self.input = self.input.saturating_add(other.input);
|
||||
self.output = self.output.saturating_add(other.output);
|
||||
self.cache_creation = self.cache_creation.saturating_add(other.cache_creation);
|
||||
self.cache_read = self.cache_read.saturating_add(other.cache_read);
|
||||
}
|
||||
}
|
||||
|
||||
/// What one transcript line contributed.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub(crate) struct RecordTally {
|
||||
pub(crate) tokens: TokenTally,
|
||||
/// The model that produced the record, when it named one.
|
||||
pub(crate) model: Option<String>,
|
||||
}
|
||||
|
||||
/// Reads one transcript line.
|
||||
///
|
||||
/// Total: a line that is not JSON, not an assistant record, or carries no
|
||||
/// usage contributes nothing. A malformed line is skipped rather than faulting
|
||||
/// the probe, because a transcript is append-only and the last line of a file
|
||||
/// being written is routinely incomplete.
|
||||
pub(crate) fn parse_record(line: &str) -> Option<RecordTally> {
|
||||
let record: RecordWire = serde_json::from_str(line).ok()?;
|
||||
if record.record_type.as_deref() != Some("assistant") {
|
||||
return None;
|
||||
}
|
||||
let message = record.message?;
|
||||
let usage = message.usage?;
|
||||
let tokens = TokenTally {
|
||||
input: usage.input,
|
||||
output: usage.output,
|
||||
cache_creation: usage.cache_creation,
|
||||
cache_read: usage.cache_read,
|
||||
};
|
||||
if tokens.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(RecordTally {
|
||||
tokens,
|
||||
model: message.model,
|
||||
})
|
||||
}
|
||||
|
||||
/// Sums the complete lines in a chunk, reporting how many bytes were consumed.
|
||||
///
|
||||
/// A trailing partial line is left unconsumed so the next read picks it up
|
||||
/// whole. This is what makes the running total monotonic: every byte is
|
||||
/// counted exactly once, and a file that is still being appended to never
|
||||
/// double-counts or loses its last record.
|
||||
pub(crate) fn consume_chunk(chunk: &str) -> (RecordTally, usize) {
|
||||
let mut tally = RecordTally::default();
|
||||
let mut consumed = 0;
|
||||
for line in chunk.split_inclusive('\n') {
|
||||
if !line.ends_with('\n') {
|
||||
break; // A partial record. Leave it for the next read.
|
||||
}
|
||||
consumed += line.len();
|
||||
if line.len() > MAX_RECORD_BYTES {
|
||||
continue;
|
||||
}
|
||||
if let Some(record) = parse_record(line.trim_end()) {
|
||||
tally.tokens.add(record.tokens);
|
||||
if record.model.is_some() {
|
||||
tally.model = record.model;
|
||||
}
|
||||
}
|
||||
}
|
||||
(tally, consumed)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{TokenTally, consume_chunk, parse_record};
|
||||
|
||||
/// Shaped exactly like a real assistant record, with the content fields a
|
||||
/// real one carries so the test proves they are ignored rather than absent.
|
||||
///
|
||||
/// Deliberately one line: a transcript is one JSON object per line, and a
|
||||
/// fixture that wrapped would not exercise the chunk boundary logic.
|
||||
const ASSISTANT: &str = concat!(
|
||||
r#"{"type":"assistant","uuid":"u1","sessionId":"s1","cwd":"/home/x","#,
|
||||
r#""message":{"id":"msg_1","role":"assistant","model":"claude-opus-5","#,
|
||||
r#""content":[{"type":"text","text":"SECRET CONVERSATION CONTENT"}],"#,
|
||||
r#""usage":{"input_tokens":10,"output_tokens":200,"#,
|
||||
r#""cache_creation_input_tokens":1000,"cache_read_input_tokens":50000,"#,
|
||||
r#""service_tier":"standard","cache_creation":{"ephemeral_5m_input_tokens":7}}}}"#,
|
||||
);
|
||||
|
||||
#[test]
|
||||
fn an_assistant_record_yields_its_four_counters() {
|
||||
let record = parse_record(ASSISTANT).expect("an assistant record with usage counts");
|
||||
assert_eq!(
|
||||
record.tokens,
|
||||
TokenTally {
|
||||
input: 10,
|
||||
output: 200,
|
||||
cache_creation: 1_000,
|
||||
cache_read: 50_000,
|
||||
}
|
||||
);
|
||||
assert_eq!(record.model.as_deref(), Some("claude-opus-5"));
|
||||
assert_eq!(record.tokens.total(), 51_210);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_content_is_not_representable() {
|
||||
// The record above carries message text. The parser's own output type
|
||||
// has nowhere to put it, so this is a property of the types rather
|
||||
// than of the parsing.
|
||||
let record = parse_record(ASSISTANT).expect("parsed");
|
||||
let rendered = format!("{record:?}");
|
||||
assert!(
|
||||
!rendered.contains("SECRET"),
|
||||
"no debug rendering of a parsed record may include message content"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_assistant_records_contribute_nothing() {
|
||||
for line in [
|
||||
r#"{"type":"user","message":{"role":"user","content":"hello"}}"#,
|
||||
r#"{"type":"attachment","attachment":{"content":"a file"}}"#,
|
||||
r#"{"type":"system","subtype":"hook"}"#,
|
||||
] {
|
||||
assert!(parse_record(line).is_none(), "{line} must not be counted");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_assistant_record_without_usage_contributes_nothing() {
|
||||
assert!(
|
||||
parse_record(r#"{"type":"assistant","message":{"model":"claude-opus-5"}}"#).is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_lines_are_skipped_rather_than_faulting() {
|
||||
for line in ["", "not json", "{", "[]", "\"a string\""] {
|
||||
assert!(parse_record(line).is_none(), "{line:?} must not panic");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_partial_trailing_line_is_left_for_the_next_read() {
|
||||
let chunk = format!("{ASSISTANT}\n{{\"type\":\"assis");
|
||||
let (tally, consumed) = consume_chunk(&chunk);
|
||||
assert_eq!(tally.tokens.total(), 51_210);
|
||||
assert_eq!(
|
||||
consumed,
|
||||
ASSISTANT.len() + 1,
|
||||
"only the complete record may be consumed"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resuming_from_the_offset_counts_every_byte_exactly_once() {
|
||||
let whole = format!("{ASSISTANT}\n{ASSISTANT}\n");
|
||||
let (first, consumed) = consume_chunk(&whole[..ASSISTANT.len() + 4]);
|
||||
let (second, _) = consume_chunk(&whole[consumed..]);
|
||||
let mut total = first.tokens;
|
||||
total.add(second.tokens);
|
||||
assert_eq!(
|
||||
total.total(),
|
||||
2 * 51_210,
|
||||
"a split read must not double-count or drop a record"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_oversized_line_is_consumed_but_not_parsed() {
|
||||
let line = format!("{}\n", "x".repeat(super::MAX_RECORD_BYTES + 1));
|
||||
let (tally, consumed) = consume_chunk(&line);
|
||||
assert_eq!(tally.tokens.total(), 0);
|
||||
assert_eq!(
|
||||
consumed,
|
||||
line.len(),
|
||||
"the offset must still advance past it"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn totals_saturate_rather_than_overflowing() {
|
||||
let tally = TokenTally {
|
||||
input: u64::MAX,
|
||||
output: u64::MAX,
|
||||
cache_creation: 0,
|
||||
cache_read: 0,
|
||||
};
|
||||
assert_eq!(tally.total(), u64::MAX);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
//! The only place a wall clock enters the usage path.
|
||||
//!
|
||||
//! [`lumbridge_core::UsageLedger::record`] rejects an observation older than
|
||||
//! the newest one already recorded for that profile, because an append-only
|
||||
//! stream cannot move backwards. A raw `SystemTime::now()` can move backwards:
|
||||
//! an NTP step, a suspend/resume, or a manual clock change would poison a
|
||||
//! profile's stream permanently. This clock anchors wall time once and then
|
||||
//! advances it with a monotonic instant, so the timestamps it emits are
|
||||
//! non-decreasing by construction while still being comparable to the Unix
|
||||
//! reset times a provider reports.
|
||||
|
||||
use std::time::{Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// A non-decreasing Unix-millisecond clock.
|
||||
#[derive(Debug)]
|
||||
pub struct MonotonicWallClock {
|
||||
anchor_unix_ms: u64,
|
||||
anchor: Instant,
|
||||
last_emitted_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for MonotonicWallClock {
|
||||
fn default() -> Self {
|
||||
Self::start()
|
||||
}
|
||||
}
|
||||
|
||||
impl MonotonicWallClock {
|
||||
/// Anchors to the system clock once, now.
|
||||
#[must_use]
|
||||
pub fn start() -> Self {
|
||||
let anchor_unix_ms = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.and_then(|elapsed| u64::try_from(elapsed.as_millis()).ok())
|
||||
.unwrap_or(0);
|
||||
Self::start_at(anchor_unix_ms, Instant::now())
|
||||
}
|
||||
|
||||
/// Anchors to a caller-supplied wall time. Exists so tests can drive the
|
||||
/// clock without waiting for real time to pass.
|
||||
#[must_use]
|
||||
pub(crate) const fn start_at(anchor_unix_ms: u64, anchor: Instant) -> Self {
|
||||
Self {
|
||||
anchor_unix_ms,
|
||||
anchor,
|
||||
last_emitted_ms: anchor_unix_ms,
|
||||
}
|
||||
}
|
||||
|
||||
/// The current time in Unix milliseconds, never earlier than the last value
|
||||
/// this clock returned.
|
||||
pub fn now_ms(&mut self) -> u64 {
|
||||
let elapsed_ms = u64::try_from(self.anchor.elapsed().as_millis()).unwrap_or(u64::MAX);
|
||||
let candidate = self.anchor_unix_ms.saturating_add(elapsed_ms);
|
||||
self.last_emitted_ms = self.last_emitted_ms.max(candidate);
|
||||
self.last_emitted_ms
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn last_emitted_ms(&self) -> u64 {
|
||||
self.last_emitted_ms
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::MonotonicWallClock;
|
||||
use std::time::Instant;
|
||||
|
||||
#[test]
|
||||
fn the_clock_starts_at_its_anchor() {
|
||||
let mut clock = MonotonicWallClock::start_at(1_700_000_000_000, Instant::now());
|
||||
assert!(clock.now_ms() >= 1_700_000_000_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emitted_times_never_move_backwards() {
|
||||
let mut clock = MonotonicWallClock::start();
|
||||
let first = clock.now_ms();
|
||||
let second = clock.now_ms();
|
||||
assert!(second >= first);
|
||||
assert_eq!(clock.last_emitted_ms(), second);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_zero_anchor_still_produces_a_usable_stream() {
|
||||
let mut clock = MonotonicWallClock::start_at(0, Instant::now());
|
||||
let first = clock.now_ms();
|
||||
assert!(clock.now_ms() >= first);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,500 @@
|
||||
//! The first real Lumbridge usage adapter: Codex `account/rateLimits/read`.
|
||||
//!
|
||||
//! Codex is the first harness because its quota surface is prose-documented
|
||||
//! upstream with a worked example, pinned by a checked-in JSON Schema, and
|
||||
//! genuinely provider-originated: the app-server forwards the backend's
|
||||
//! `x-codex-*-used-percent` response headers rather than computing a number.
|
||||
//! Lumbridge obtains it by launching `codex app-server` and letting the harness
|
||||
//! resolve its own credentials from `CODEX_HOME`. Lumbridge never reads an auth
|
||||
//! file, and [`crate::jsonrpc`] makes asking for a token unrepresentable.
|
||||
//!
|
||||
//! Three threads, each with one job. The caller's thread starts the child, so a
|
||||
//! missing harness is an immediate error rather than a silent fault. A reader
|
||||
//! thread does the blocking pipe reads. A worker thread owns the protocol clock
|
||||
//! and the writer.
|
||||
//!
|
||||
//! Shutdown is bounded and never waits on the far end. It kills the child's
|
||||
//! whole process group, sets a stop flag, and joins only the worker, which
|
||||
//! checks that flag every tick. The reader is deliberately never joined: it is
|
||||
//! blocked in a pipe read, and while killing the group is what normally closes
|
||||
//! that pipe, a descendant that escaped the group by calling `setsid` would
|
||||
//! otherwise hold shutdown open forever. A leaked thread parked on a dead file
|
||||
//! descriptor is a far better failure than a hung UI, since `Drop` calls
|
||||
//! shutdown.
|
||||
|
||||
mod parse;
|
||||
mod session;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::mpsc::{self, Receiver, RecvTimeoutError, SyncSender, TryRecvError, TrySendError};
|
||||
use std::thread::{self, JoinHandle};
|
||||
use std::time::Duration;
|
||||
|
||||
use lumbridge_core::{AccountProfile, UsageObservation};
|
||||
|
||||
use crate::HarnessError;
|
||||
use crate::child::{ProbeChild, ProbeKiller, ProbeReader, ProbeWriter};
|
||||
use crate::clock::MonotonicWallClock;
|
||||
use crate::jsonrpc::ServerRequestClass;
|
||||
use crate::probe::{ProbeHealth, ProbeOutcome, UsageProbe};
|
||||
|
||||
use session::{CodexSession, SessionStep};
|
||||
|
||||
/// The shortest interval a probe may poll an account quota at. A quota window
|
||||
/// is measured in minutes; polling faster only burns the user's process table.
|
||||
const MIN_POLL_INTERVAL: Duration = Duration::from_secs(1);
|
||||
/// How long the worker waits for a line before re-checking its poll timer.
|
||||
const WORKER_TICK: Duration = Duration::from_millis(100);
|
||||
const EVENT_QUEUE: usize = 64;
|
||||
const LINE_QUEUE: usize = 64;
|
||||
|
||||
/// How the probe is launched.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct CodexProbeOptions {
|
||||
/// The program to run. Defaults to `codex` on `PATH`.
|
||||
pub program: String,
|
||||
/// How often to request a fresh reading.
|
||||
pub poll_interval: Duration,
|
||||
/// Reported to the harness during the handshake. Not a credential.
|
||||
pub client_name: String,
|
||||
pub client_version: String,
|
||||
}
|
||||
|
||||
impl Default for CodexProbeOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
program: "codex".to_owned(),
|
||||
poll_interval: Duration::from_secs(60),
|
||||
client_name: "lumbridge".to_owned(),
|
||||
client_version: env!("CARGO_PKG_VERSION").to_owned(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Replays a synthetic app-server transcript through the real session logic.
|
||||
///
|
||||
/// This is the same state machine the live probe runs, with the process, the
|
||||
/// socket, and the clock removed: each frame carries the moment it is to be
|
||||
/// treated as observed. It exists so an adapter can be verified end to end
|
||||
/// against a fixture, and so a captured (synthetic) transcript can be checked
|
||||
/// offline when upstream changes its wire shape.
|
||||
///
|
||||
/// Frames the session would answer are answered internally and dropped; only
|
||||
/// the observations it decided to record come back.
|
||||
#[must_use]
|
||||
pub fn replay_transcript(frames: &[(&str, u64)]) -> Vec<UsageObservation> {
|
||||
let Ok(mut session) = CodexSession::new() else {
|
||||
return Vec::new();
|
||||
};
|
||||
let _ = session.initialize("lumbridge-replay", "0");
|
||||
let mut observations = Vec::new();
|
||||
for (line, observed_at_ms) in frames {
|
||||
// A replay polls whenever the session will let it, so a transcript
|
||||
// does not have to encode the worker's timer.
|
||||
let _ = session.poll_frame();
|
||||
match session.on_line(line, *observed_at_ms) {
|
||||
SessionStep::Observations(mut batch) => observations.append(&mut batch),
|
||||
SessionStep::Idle
|
||||
| SessionStep::Send(_)
|
||||
| SessionStep::Health(_)
|
||||
| SessionStep::Refused { .. } => {}
|
||||
}
|
||||
}
|
||||
observations
|
||||
}
|
||||
|
||||
/// What the worker needs to run the handshake and its poll timer.
|
||||
struct WorkerSetup {
|
||||
client_name: String,
|
||||
client_version: String,
|
||||
poll_interval_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum WorkerEvent {
|
||||
Observations(Vec<UsageObservation>),
|
||||
Health(ProbeHealth),
|
||||
RefusedServerRequest(ServerRequestClass),
|
||||
}
|
||||
|
||||
/// A running Codex quota probe.
|
||||
pub struct CodexProbe {
|
||||
profiles: Vec<AccountProfile>,
|
||||
events: Option<Receiver<WorkerEvent>>,
|
||||
killer: ProbeKiller,
|
||||
stop: Arc<AtomicBool>,
|
||||
worker: Option<JoinHandle<()>>,
|
||||
health: ProbeHealth,
|
||||
refused_credential_requests: u64,
|
||||
}
|
||||
|
||||
impl CodexProbe {
|
||||
/// Starts the probe, its reader, and its worker.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`HarnessError::ProbeIntervalTooShort`] below one second,
|
||||
/// [`HarnessError::EmptyProgram`] for a blank program name, and
|
||||
/// [`HarnessError::SpawnFailed`] when the harness cannot be launched.
|
||||
pub fn start(options: CodexProbeOptions) -> Result<Self, HarnessError> {
|
||||
if options.poll_interval < MIN_POLL_INTERVAL {
|
||||
return Err(HarnessError::ProbeIntervalTooShort);
|
||||
}
|
||||
let CodexProbeOptions {
|
||||
program,
|
||||
poll_interval,
|
||||
client_name,
|
||||
client_version,
|
||||
} = options;
|
||||
|
||||
let session = CodexSession::new()?;
|
||||
let profiles = session.profiles().to_vec();
|
||||
let (writer, reader, killer) = ProbeChild::spawn(&program, &["app-server"])?.into_parts();
|
||||
|
||||
let (line_sender, lines) = mpsc::sync_channel(LINE_QUEUE);
|
||||
let (event_sender, events) = mpsc::sync_channel(EVENT_QUEUE);
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
|
||||
// Deliberately not retained: the reader is never joined. See the module
|
||||
// documentation for why waiting on a blocking pipe read is unsafe here.
|
||||
thread::Builder::new()
|
||||
.name("lumbridge-codex-probe-reader".to_owned())
|
||||
.spawn(move || run_reader(reader, &line_sender))
|
||||
.map_err(|error| HarnessError::SpawnFailed(error.kind()))?;
|
||||
|
||||
let setup = WorkerSetup {
|
||||
client_name,
|
||||
client_version,
|
||||
poll_interval_ms: u64::try_from(poll_interval.as_millis()).unwrap_or(u64::MAX),
|
||||
};
|
||||
let worker_stop = Arc::clone(&stop);
|
||||
let worker = thread::Builder::new()
|
||||
.name("lumbridge-codex-probe".to_owned())
|
||||
.spawn(move || {
|
||||
run_worker(session, &setup, writer, &lines, &event_sender, &worker_stop);
|
||||
})
|
||||
.map_err(|error| HarnessError::SpawnFailed(error.kind()))?;
|
||||
|
||||
Ok(Self {
|
||||
profiles,
|
||||
events: Some(events),
|
||||
killer,
|
||||
stop,
|
||||
worker: Some(worker),
|
||||
health: ProbeHealth::Starting,
|
||||
refused_credential_requests: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// How many times the harness asked this probe for a credential and was
|
||||
/// refused. A non-zero count is not an error; it is evidence that the
|
||||
/// refusal path ran.
|
||||
#[must_use]
|
||||
pub const fn refused_credential_requests(&self) -> u64 {
|
||||
self.refused_credential_requests
|
||||
}
|
||||
}
|
||||
|
||||
impl UsageProbe for CodexProbe {
|
||||
fn poll(&mut self) -> ProbeOutcome {
|
||||
let Some(events) = self.events.as_ref() else {
|
||||
return ProbeOutcome::idle(self.health);
|
||||
};
|
||||
let mut observations = Vec::new();
|
||||
loop {
|
||||
match events.try_recv() {
|
||||
Ok(WorkerEvent::Observations(mut batch)) => observations.append(&mut batch),
|
||||
Ok(WorkerEvent::Health(health)) => self.health = health,
|
||||
Ok(WorkerEvent::RefusedServerRequest(class)) => {
|
||||
if class == ServerRequestClass::CredentialRefresh {
|
||||
self.refused_credential_requests =
|
||||
self.refused_credential_requests.saturating_add(1);
|
||||
}
|
||||
}
|
||||
Err(TryRecvError::Empty) => break,
|
||||
Err(TryRecvError::Disconnected) => {
|
||||
if !self.health.is_faulted() {
|
||||
self.health = ProbeHealth::Stopped;
|
||||
}
|
||||
self.events = None;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
ProbeOutcome::new(observations, self.health)
|
||||
}
|
||||
|
||||
fn profiles(&self) -> &[AccountProfile] {
|
||||
&self.profiles
|
||||
}
|
||||
|
||||
fn shutdown(&mut self) {
|
||||
// Order matters. Signal first so the worker cannot start another poll,
|
||||
// then kill the whole group, then join only the worker: it wakes at
|
||||
// most one tick later regardless of what the child is doing.
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
self.killer.kill();
|
||||
self.events = None;
|
||||
if let Some(worker) = self.worker.take() {
|
||||
let _ = worker.join();
|
||||
}
|
||||
if !self.health.is_faulted() {
|
||||
self.health = ProbeHealth::Stopped;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CodexProbe {
|
||||
fn drop(&mut self) {
|
||||
self.shutdown();
|
||||
}
|
||||
}
|
||||
|
||||
/// Blocking reads, and nothing else.
|
||||
fn run_reader(mut reader: ProbeReader, lines: &SyncSender<Result<String, HarnessError>>) {
|
||||
loop {
|
||||
let outcome = reader.read_line();
|
||||
let terminal = outcome.is_err();
|
||||
if lines.send(outcome).is_err() || terminal {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Owns the protocol clock and the writer. Holds no policy of its own.
|
||||
fn run_worker(
|
||||
mut session: CodexSession,
|
||||
setup: &WorkerSetup,
|
||||
mut writer: ProbeWriter,
|
||||
lines: &Receiver<Result<String, HarnessError>>,
|
||||
events: &SyncSender<WorkerEvent>,
|
||||
stop: &AtomicBool,
|
||||
) {
|
||||
let mut clock = MonotonicWallClock::start();
|
||||
let handshake = session.initialize(&setup.client_name, &setup.client_version);
|
||||
if let Err(error) = writer.write_line(&handshake) {
|
||||
let _ = events.try_send(WorkerEvent::Health(ProbeHealth::Faulted(error)));
|
||||
return;
|
||||
}
|
||||
|
||||
// Zero means "poll as soon as the handshake lets us", because
|
||||
// `poll_frame` returns nothing until the session is ready.
|
||||
let mut last_poll_ms = 0;
|
||||
loop {
|
||||
if stop.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
let now_ms = clock.now_ms();
|
||||
if now_ms.saturating_sub(last_poll_ms) >= setup.poll_interval_ms
|
||||
&& let Some(frame) = session.poll_frame()
|
||||
{
|
||||
last_poll_ms = now_ms;
|
||||
if writer.write_line(&frame).is_err() {
|
||||
let _ =
|
||||
events.try_send(WorkerEvent::Health(ProbeHealth::Faulted(HarnessError::Eof)));
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
match lines.recv_timeout(WORKER_TICK) {
|
||||
Ok(Ok(line)) => {
|
||||
if !dispatch(&mut session, &mut writer, events, &line, clock.now_ms()) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Err(error)) => {
|
||||
let _ = events.try_send(WorkerEvent::Health(ProbeHealth::Faulted(error)));
|
||||
break;
|
||||
}
|
||||
Err(RecvTimeoutError::Timeout) => {}
|
||||
Err(RecvTimeoutError::Disconnected) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Applies one line. Returns whether the worker should keep running.
|
||||
fn dispatch(
|
||||
session: &mut CodexSession,
|
||||
writer: &mut ProbeWriter,
|
||||
events: &SyncSender<WorkerEvent>,
|
||||
line: &str,
|
||||
now_ms: u64,
|
||||
) -> bool {
|
||||
match session.on_line(line.trim_end(), now_ms) {
|
||||
SessionStep::Idle => true,
|
||||
SessionStep::Send(frame) => {
|
||||
if writer.write_line(&frame).is_err() {
|
||||
return false;
|
||||
}
|
||||
if session.is_ready() {
|
||||
let _ = events.try_send(WorkerEvent::Health(ProbeHealth::Ready));
|
||||
}
|
||||
true
|
||||
}
|
||||
SessionStep::Refused { line, class } => {
|
||||
let _ = events.try_send(WorkerEvent::RefusedServerRequest(class));
|
||||
writer.write_line(&line).is_ok()
|
||||
}
|
||||
SessionStep::Observations(observations) => !matches!(
|
||||
events.try_send(WorkerEvent::Observations(observations)),
|
||||
Err(TrySendError::Disconnected(_))
|
||||
),
|
||||
SessionStep::Health(health) => {
|
||||
let keep_running = !health.is_faulted();
|
||||
let _ = events.try_send(WorkerEvent::Health(health));
|
||||
keep_running
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{CodexProbe, CodexProbeOptions};
|
||||
use crate::HarnessError;
|
||||
use crate::probe::{ProbeHealth, UsageProbe};
|
||||
use std::time::Duration;
|
||||
|
||||
/// A synthetic launcher that ignores its arguments, backgrounds a helper
|
||||
/// that inherits stdio, and waits. This is the shape of a wrapper script
|
||||
/// that execs the real program as a grandchild, which is how the npm
|
||||
/// distribution of Codex behaves.
|
||||
fn launcher_script(tag: &str) -> Option<std::path::PathBuf> {
|
||||
use std::io::Write;
|
||||
|
||||
let path = std::env::temp_dir().join(format!(
|
||||
"lumbridge-probe-launcher-{}-{tag}.sh",
|
||||
std::process::id()
|
||||
));
|
||||
let mut file = std::fs::File::create(&path).ok()?;
|
||||
file.write_all(b"#!/bin/sh\nsleep 600 &\nwait\n").ok()?;
|
||||
drop(file);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o755)).ok()?;
|
||||
}
|
||||
Some(path)
|
||||
}
|
||||
|
||||
fn options(program: &str) -> CodexProbeOptions {
|
||||
CodexProbeOptions {
|
||||
program: program.to_owned(),
|
||||
poll_interval: Duration::from_secs(1),
|
||||
..CodexProbeOptions::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_too_frequent_poll_interval_is_refused() {
|
||||
let options = CodexProbeOptions {
|
||||
poll_interval: Duration::from_millis(10),
|
||||
..CodexProbeOptions::default()
|
||||
};
|
||||
assert_eq!(
|
||||
CodexProbe::start(options).err(),
|
||||
Some(HarnessError::ProbeIntervalTooShort)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_empty_program_is_refused() {
|
||||
assert_eq!(
|
||||
CodexProbe::start(options(" ")).err(),
|
||||
Some(HarnessError::EmptyProgram)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_missing_harness_fails_at_start_rather_than_silently() {
|
||||
assert!(matches!(
|
||||
CodexProbe::start(options("lumbridge-no-such-harness")).err(),
|
||||
Some(HarnessError::SpawnFailed(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_silent_harness_can_still_be_shut_down() {
|
||||
let Some(script) = launcher_script("silent") else {
|
||||
return;
|
||||
};
|
||||
let Some(program) = script.to_str() else {
|
||||
return;
|
||||
};
|
||||
let Ok(mut probe) = CodexProbe::start(options(program)) else {
|
||||
let _ = std::fs::remove_file(&script);
|
||||
return;
|
||||
};
|
||||
assert_eq!(probe.profiles().len(), 2);
|
||||
assert_eq!(probe.poll().health(), ProbeHealth::Starting);
|
||||
probe.shutdown();
|
||||
assert!(matches!(
|
||||
probe.poll().health(),
|
||||
ProbeHealth::Stopped | ProbeHealth::Faulted(_)
|
||||
));
|
||||
let _ = std::fs::remove_file(&script);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shutdown_is_idempotent() {
|
||||
let Some(script) = launcher_script("idempotent") else {
|
||||
return;
|
||||
};
|
||||
let Some(program) = script.to_str() else {
|
||||
return;
|
||||
};
|
||||
let Ok(mut probe) = CodexProbe::start(options(program)) else {
|
||||
let _ = std::fs::remove_file(&script);
|
||||
return;
|
||||
};
|
||||
probe.shutdown();
|
||||
probe.shutdown();
|
||||
let _ = std::fs::remove_file(&script);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shutdown_returns_promptly_behind_a_launcher_that_never_speaks() {
|
||||
// The grandchild inherits our stdout and never writes, so the reader is
|
||||
// blocked in a pipe read that will never complete on its own. Before
|
||||
// the process-group kill and the bounded join, shutdown blocked here
|
||||
// until the orphan exited — ten minutes, in this fixture.
|
||||
let Some(script) = launcher_script("prompt") else {
|
||||
return;
|
||||
};
|
||||
let Some(program) = script.to_str() else {
|
||||
return;
|
||||
};
|
||||
let Ok(mut probe) = CodexProbe::start(options(program)) else {
|
||||
let _ = std::fs::remove_file(&script);
|
||||
return;
|
||||
};
|
||||
let started = std::time::Instant::now();
|
||||
probe.shutdown();
|
||||
let elapsed = started.elapsed();
|
||||
let _ = std::fs::remove_file(&script);
|
||||
assert!(
|
||||
elapsed < Duration::from_secs(5),
|
||||
"shutdown must not wait on the far end, took {elapsed:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_harness_that_exits_immediately_faults_the_probe() {
|
||||
let Ok(mut probe) = CodexProbe::start(options("true")) else {
|
||||
return;
|
||||
};
|
||||
let mut health = probe.poll().health();
|
||||
for _ in 0..200 {
|
||||
if health != ProbeHealth::Starting {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
health = probe.poll().health();
|
||||
}
|
||||
assert!(
|
||||
matches!(health, ProbeHealth::Faulted(_)),
|
||||
"a harness that closes its output must surface as a fault, got {health:?}"
|
||||
);
|
||||
probe.shutdown();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,397 @@
|
||||
//! The Codex rate-limit parser: hostile input in, a closed decision out.
|
||||
//!
|
||||
//! Every skew, roll, overflow, and out-of-range case is decided here, by pure
|
||||
//! functions over deserialized wire values with no process, no clock, and no
|
||||
//! network. That is what makes the honesty rules testable.
|
||||
//!
|
||||
//! Wire shape adapted from the openai/codex app-server protocol (Apache-2.0),
|
||||
//! commit `17e8101699c5062117d0d37f504313e8af53b043`:
|
||||
//! `codex-rs/app-server/README.md` section "7) Rate limits (`ChatGPT`)" and the
|
||||
//! checked-in schema `codex_app_server_protocol.v2.schemas.json`, where
|
||||
//! `RateLimitWindow` requires only `usedPercent` and both
|
||||
//! `windowDurationMins` and `resetsAt` are nullable. No upstream code is
|
||||
//! copied; only the observable wire contract is mirrored.
|
||||
|
||||
use lumbridge_core::{AccountProfileId, UsageObservation, UsageProvenance, UsageUnit, UsageWindow};
|
||||
use serde::Deserialize;
|
||||
|
||||
/// Percent is an integer 0..=100 on the wire, so the honest resolution of a
|
||||
/// Codex reading is one percent even though the ledger stores permille.
|
||||
const MAX_USED_PERCENT: i64 = 100;
|
||||
const PERMILLE_PER_PERCENT: u64 = 10;
|
||||
const MILLIS_PER_SECOND: i64 = 1_000;
|
||||
const MILLIS_PER_MINUTE: i64 = 60_000;
|
||||
/// A permille ceiling makes the window fraction a bounded quantity, which is
|
||||
/// what lets the ledger derive an exhaustion estimate. It is not an invented
|
||||
/// quota: a percentage is by definition a share of one whole.
|
||||
const PERMILLE_CEILING: u64 = 1_000;
|
||||
|
||||
/// One quota window exactly as Codex sends it.
|
||||
///
|
||||
/// Only the three documented fields are deserialized. Sibling fields such as
|
||||
/// `rateLimitResetCredits`, `individualLimit`, `accountId`, and
|
||||
/// `rateLimitUpsell` are deliberately not modelled: a probe should not be able
|
||||
/// to accidentally read account metadata it has no use for.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct RateLimitWindowWire {
|
||||
/// Widened to `i64` so an out-of-range wire value reaches our own check
|
||||
/// rather than failing inside serde with a less specific error.
|
||||
pub(crate) used_percent: i64,
|
||||
#[serde(default)]
|
||||
pub(crate) window_duration_mins: Option<i64>,
|
||||
#[serde(default)]
|
||||
pub(crate) resets_at: Option<i64>,
|
||||
}
|
||||
|
||||
/// The snapshot. Every field is optional in the upstream schema, including
|
||||
/// `primary`, so absence must be a first-class case rather than a parse error.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct RateLimitSnapshotWire {
|
||||
#[serde(default)]
|
||||
pub(crate) primary: Option<RateLimitWindowWire>,
|
||||
#[serde(default)]
|
||||
pub(crate) secondary: Option<RateLimitWindowWire>,
|
||||
}
|
||||
|
||||
/// The `account/rateLimits/read` result body, which is also the shape of the
|
||||
/// `account/rateLimits/updated` notification's params.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct RateLimitsResultWire {
|
||||
#[serde(default)]
|
||||
pub(crate) rate_limits: RateLimitSnapshotWire,
|
||||
}
|
||||
|
||||
/// Why a syntactically valid window still cannot be recorded as a fact.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum UnusableReason {
|
||||
/// A percentage outside 0..=100. The relay is not telling the truth.
|
||||
PercentOutOfRange,
|
||||
/// The window had already ended when the reading was taken. Decision 0012
|
||||
/// requires an expired window to read as rolled over rather than as a
|
||||
/// frozen percentage from a window that is gone.
|
||||
WindowEnded,
|
||||
/// A reset timestamp that cannot be represented as Unix milliseconds.
|
||||
ResetsAtOutOfRange,
|
||||
}
|
||||
|
||||
/// What the parser decided about one window.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum WindowReading {
|
||||
/// A usable share of the window, with a window attached when the wire
|
||||
/// carried enough to build one.
|
||||
Usable {
|
||||
permille: u64,
|
||||
window: Option<UsageWindow>,
|
||||
},
|
||||
/// The reading must be recorded as an explicit gap.
|
||||
Unusable(UnusableReason),
|
||||
}
|
||||
|
||||
/// Decides what a single Codex window means at a point in time.
|
||||
///
|
||||
/// Total: every input produces a decision, and none of them panics.
|
||||
pub(crate) fn parse_window(wire: &RateLimitWindowWire, observed_at_ms: u64) -> WindowReading {
|
||||
if wire.used_percent < 0 || wire.used_percent > MAX_USED_PERCENT {
|
||||
return WindowReading::Unusable(UnusableReason::PercentOutOfRange);
|
||||
}
|
||||
// Lossless: used_percent is already proven to be 0..=100.
|
||||
let permille = u64::try_from(wire.used_percent).unwrap_or(0) * PERMILLE_PER_PERCENT;
|
||||
|
||||
let Some(resets_at_seconds) = wire.resets_at else {
|
||||
// A percentage with no reset time is still a true percentage. It just
|
||||
// cannot answer "when does this refill?".
|
||||
return WindowReading::Usable {
|
||||
permille,
|
||||
window: None,
|
||||
};
|
||||
};
|
||||
let Some(resets_at_ms) = unix_seconds_to_millis(resets_at_seconds) else {
|
||||
return WindowReading::Unusable(UnusableReason::ResetsAtOutOfRange);
|
||||
};
|
||||
if observed_at_ms >= resets_at_ms {
|
||||
return WindowReading::Unusable(UnusableReason::WindowEnded);
|
||||
}
|
||||
|
||||
// The reset is a fact on its own. Codex declares the duration and the reset
|
||||
// independently, so a missing or unreconcilable start must cost the start
|
||||
// and nothing else: the reset still bounds the exhaustion forecast and
|
||||
// still identifies which quota period this reading belongs to.
|
||||
let window = wire
|
||||
.window_duration_mins
|
||||
.and_then(|minutes| window_start_ms(resets_at_ms, minutes))
|
||||
// A start after the observation means the local clock disagrees with
|
||||
// the provider's. That disagreement is about the start, not the reset.
|
||||
.filter(|started_at_ms| *started_at_ms <= observed_at_ms)
|
||||
.and_then(|started_at_ms| UsageWindow::new(started_at_ms, resets_at_ms).ok())
|
||||
.unwrap_or_else(|| UsageWindow::until(resets_at_ms));
|
||||
|
||||
WindowReading::Usable {
|
||||
permille,
|
||||
window: Some(window),
|
||||
}
|
||||
}
|
||||
|
||||
fn unix_seconds_to_millis(seconds: i64) -> Option<u64> {
|
||||
let millis = seconds.checked_mul(MILLIS_PER_SECOND)?;
|
||||
u64::try_from(millis).ok()
|
||||
}
|
||||
|
||||
fn window_start_ms(resets_at_ms: u64, duration_minutes: i64) -> Option<u64> {
|
||||
if duration_minutes <= 0 {
|
||||
return None;
|
||||
}
|
||||
let duration_ms = u64::try_from(duration_minutes.checked_mul(MILLIS_PER_MINUTE)?).ok()?;
|
||||
resets_at_ms.checked_sub(duration_ms)
|
||||
}
|
||||
|
||||
/// Turns a decision into the observation the ledger should hold.
|
||||
///
|
||||
/// A usable reading is `ProviderReported`: Codex forwards the backend's
|
||||
/// `x-codex-*-used-percent` response headers rather than computing the value,
|
||||
/// so the provider is the one making the claim. Anything else is an explicit
|
||||
/// gap, never a zero.
|
||||
pub(crate) fn to_observation(
|
||||
profile: &AccountProfileId,
|
||||
reading: WindowReading,
|
||||
observed_at_ms: u64,
|
||||
) -> UsageObservation {
|
||||
let WindowReading::Usable { permille, window } = reading else {
|
||||
return UsageObservation::unavailable(profile.clone(), observed_at_ms);
|
||||
};
|
||||
let Ok(observation) = UsageObservation::counted(
|
||||
profile.clone(),
|
||||
UsageUnit::WindowPermille,
|
||||
permille,
|
||||
UsageProvenance::ProviderReported,
|
||||
observed_at_ms,
|
||||
) else {
|
||||
return UsageObservation::unavailable(profile.clone(), observed_at_ms);
|
||||
};
|
||||
let observation = observation
|
||||
.with_limit(PERMILLE_CEILING)
|
||||
.unwrap_or_else(|_| UsageObservation::unavailable(profile.clone(), observed_at_ms));
|
||||
let Some(window) = window else {
|
||||
return observation;
|
||||
};
|
||||
// A window that no longer contains the observation is rejected by core.
|
||||
// Keeping the windowless observation preserves the provider's percentage
|
||||
// instead of discarding a true fact over a timing disagreement.
|
||||
observation
|
||||
.clone()
|
||||
.with_window(window)
|
||||
.unwrap_or(observation)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
RateLimitSnapshotWire, RateLimitWindowWire, RateLimitsResultWire, UnusableReason,
|
||||
WindowReading, parse_window, to_observation,
|
||||
};
|
||||
use lumbridge_core::{AccountProfileId, UsageProvenance, UsageUnit, UsageWindow};
|
||||
|
||||
const HOUR_MS: u64 = 3_600_000;
|
||||
const RESETS_AT_SECONDS: i64 = 1_730_947_200;
|
||||
const RESETS_AT_MS: u64 = 1_730_947_200_000;
|
||||
|
||||
fn profile() -> AccountProfileId {
|
||||
AccountProfileId::new("codex-app-server-primary").expect("a valid fixture ID")
|
||||
}
|
||||
|
||||
fn wire(
|
||||
used_percent: i64,
|
||||
minutes: Option<i64>,
|
||||
resets_at: Option<i64>,
|
||||
) -> RateLimitWindowWire {
|
||||
RateLimitWindowWire {
|
||||
used_percent,
|
||||
window_duration_mins: minutes,
|
||||
resets_at,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_documented_example_parses_to_a_bounded_window() {
|
||||
// The worked example from the upstream README, read one hour before it
|
||||
// resets: { "usedPercent": 25, "windowDurationMins": 15, ... } widened
|
||||
// to a duration long enough to contain the observation.
|
||||
let reading = parse_window(
|
||||
&wire(25, Some(300), Some(RESETS_AT_SECONDS)),
|
||||
RESETS_AT_MS - HOUR_MS,
|
||||
);
|
||||
let WindowReading::Usable { permille, window } = reading else {
|
||||
panic!("the documented example must be usable");
|
||||
};
|
||||
assert_eq!(permille, 250);
|
||||
let window = window.expect("a duration and a reset build a window");
|
||||
assert_eq!(window.resets_at_ms(), RESETS_AT_MS);
|
||||
assert_eq!(window.started_at_ms(), Some(RESETS_AT_MS - 300 * 60_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_reading_taken_after_the_reset_is_a_rolled_window_not_a_frozen_percentage() {
|
||||
let reading = parse_window(
|
||||
&wire(93, Some(300), Some(RESETS_AT_SECONDS)),
|
||||
RESETS_AT_MS + 1,
|
||||
);
|
||||
assert_eq!(
|
||||
reading,
|
||||
WindowReading::Unusable(UnusableReason::WindowEnded)
|
||||
);
|
||||
let observation = to_observation(&profile(), reading, RESETS_AT_MS + 1);
|
||||
assert_eq!(observation.provenance(), UsageProvenance::Unavailable);
|
||||
assert_eq!(observation.consumed(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_absent_duration_keeps_the_percentage_and_the_reset() {
|
||||
// Codex declares the duration and the reset independently. Dropping
|
||||
// the reset for want of a duration would discard a fact the provider
|
||||
// stated, and would let a burn baseline be drawn across two quota
|
||||
// periods while the exhaustion guard never fired.
|
||||
let reading = parse_window(
|
||||
&wire(40, None, Some(RESETS_AT_SECONDS)),
|
||||
RESETS_AT_MS - HOUR_MS,
|
||||
);
|
||||
assert_eq!(
|
||||
reading,
|
||||
WindowReading::Usable {
|
||||
permille: 400,
|
||||
window: Some(UsageWindow::until(RESETS_AT_MS))
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_absent_reset_keeps_the_percentage_and_drops_the_window() {
|
||||
let reading = parse_window(&wire(40, Some(300), None), 5_000);
|
||||
assert_eq!(
|
||||
reading,
|
||||
WindowReading::Usable {
|
||||
permille: 400,
|
||||
window: None
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_start_after_the_observation_costs_the_start_and_not_the_reset() {
|
||||
// A 15 minute window resetting far in the future starts after now:
|
||||
// the local clock and the provider disagree about the start only.
|
||||
let reading = parse_window(&wire(25, Some(15), Some(RESETS_AT_SECONDS)), 1_000);
|
||||
assert_eq!(
|
||||
reading,
|
||||
WindowReading::Usable {
|
||||
permille: 250,
|
||||
window: Some(UsageWindow::until(RESETS_AT_MS))
|
||||
},
|
||||
"clock disagreement must not discard the provider's reset time"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_out_of_range_percentage_is_refused() {
|
||||
for percent in [-1, 101, i64::MAX, i64::MIN] {
|
||||
assert_eq!(
|
||||
parse_window(&wire(percent, Some(300), Some(RESETS_AT_SECONDS)), 0),
|
||||
WindowReading::Unusable(UnusableReason::PercentOutOfRange),
|
||||
"percent {percent} must not be recorded"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unrepresentable_reset_is_refused_without_overflowing() {
|
||||
for seconds in [-1, i64::MIN, i64::MAX] {
|
||||
assert!(matches!(
|
||||
parse_window(&wire(25, Some(300), Some(seconds)), 0),
|
||||
WindowReading::Unusable(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_zero_or_negative_duration_does_not_underflow() {
|
||||
for minutes in [0, -5, i64::MIN] {
|
||||
let reading = parse_window(&wire(25, Some(minutes), Some(RESETS_AT_SECONDS)), 1_000);
|
||||
assert_eq!(
|
||||
reading,
|
||||
WindowReading::Usable {
|
||||
permille: 250,
|
||||
window: Some(UsageWindow::until(RESETS_AT_MS))
|
||||
},
|
||||
"a nonsensical duration must not cost the reset"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_weekly_window_matches_the_shape_a_live_account_returns() {
|
||||
// Observed from a real account: usedPercent 18, windowDurationMins
|
||||
// 10080 (one week), resetsAt in Unix seconds, secondary null.
|
||||
let week_minutes: i64 = 10_080;
|
||||
let observed_at_ms = RESETS_AT_MS - 24 * HOUR_MS;
|
||||
let reading = parse_window(
|
||||
&wire(18, Some(week_minutes), Some(RESETS_AT_SECONDS)),
|
||||
observed_at_ms,
|
||||
);
|
||||
let WindowReading::Usable { permille, window } = reading else {
|
||||
panic!("a weekly window is usable");
|
||||
};
|
||||
assert_eq!(permille, 180);
|
||||
let window = window.expect("a week-long window is still a window");
|
||||
assert_eq!(
|
||||
window.started_at_ms(),
|
||||
Some(RESETS_AT_MS - u64::try_from(week_minutes).expect("a positive duration") * 60_000)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_usable_reading_becomes_a_provider_reported_permille_observation() {
|
||||
let reading = parse_window(
|
||||
&wire(25, Some(300), Some(RESETS_AT_SECONDS)),
|
||||
RESETS_AT_MS - HOUR_MS,
|
||||
);
|
||||
let observation = to_observation(&profile(), reading, RESETS_AT_MS - HOUR_MS);
|
||||
assert_eq!(observation.provenance(), UsageProvenance::ProviderReported);
|
||||
assert_eq!(observation.unit(), Some(UsageUnit::WindowPermille));
|
||||
assert_eq!(observation.consumed(), 250);
|
||||
assert_eq!(observation.limit(), Some(1_000));
|
||||
assert!(observation.window().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_absent_primary_window_deserializes_rather_than_failing() {
|
||||
let body: RateLimitsResultWire = serde_json::from_str(
|
||||
r#"{"rateLimits":{"primary":null,"secondary":null,"rateLimitReachedType":null}}"#,
|
||||
)
|
||||
.expect("every snapshot field is optional upstream");
|
||||
assert!(body.rate_limits.primary.is_none());
|
||||
assert!(body.rate_limits.secondary.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unmodelled_sibling_fields_are_ignored_not_rejected() {
|
||||
// The upstream result also carries rateLimitResetCredits, accountId,
|
||||
// individualLimit and rateLimitUpsell. A probe must tolerate them and
|
||||
// must not model them.
|
||||
let body: RateLimitsResultWire = serde_json::from_str(
|
||||
r#"{"rateLimits":{"primary":{"usedPercent":25,"windowDurationMins":15,
|
||||
"resetsAt":1730947200},"individualLimit":null,"accountId":"acct_x",
|
||||
"planType":"plus"},"rateLimitResetCredits":{"availableCount":2}}"#,
|
||||
)
|
||||
.expect("unknown fields are ignored");
|
||||
let primary = body.rate_limits.primary.expect("primary is present");
|
||||
assert_eq!(primary.used_percent, 25);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_snapshot_with_no_windows_defaults_rather_than_panicking() {
|
||||
let snapshot = RateLimitSnapshotWire::default();
|
||||
assert!(snapshot.primary.is_none());
|
||||
assert!(snapshot.secondary.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,584 @@
|
||||
//! The Codex probe's protocol logic, as a pure state machine.
|
||||
//!
|
||||
//! Every decision the probe makes about what to send, what to record, and what
|
||||
//! to refuse happens here, driven by `&str` lines and an explicit timestamp.
|
||||
//! No process, no socket, and no clock are involved, so the whole protocol is
|
||||
//! testable from string fixtures. The thread in [`super::worker`] does nothing
|
||||
//! but move bytes between a child process and this type.
|
||||
|
||||
use lumbridge_core::{AccountProfile, AccountProfileId, UsageObservation, UsageWindow};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::HarnessError;
|
||||
use crate::jsonrpc::{
|
||||
AcceptedNotification, INVALID_REQUEST, Inbound, OutboundRequest, ServerRequestClass, classify,
|
||||
encode_notification, encode_refusal, encode_request,
|
||||
};
|
||||
use crate::probe::ProbeHealth;
|
||||
|
||||
use super::parse::{
|
||||
RateLimitSnapshotWire, RateLimitsResultWire, WindowReading, parse_window, to_observation,
|
||||
};
|
||||
|
||||
/// Re-emit an unchanged reading at least this often so the ledger keeps a
|
||||
/// baseline for its burn rate. Without a floor, a quiet account would hold one
|
||||
/// lone observation and the footer could never derive a rate; without a
|
||||
/// ceiling on how often we re-record, the bounded retention would be flushed
|
||||
/// of real history by identical rows.
|
||||
const MIN_REEMIT_INTERVAL_MS: u64 = 300_000;
|
||||
|
||||
/// The two windows Codex reports are two separate facts about two separate
|
||||
/// quota periods. Merging them would mean silently choosing one.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum WindowSlot {
|
||||
Primary,
|
||||
Secondary,
|
||||
}
|
||||
|
||||
impl WindowSlot {
|
||||
const ALL: [Self; 2] = [Self::Primary, Self::Secondary];
|
||||
|
||||
/// Static, account-free identifiers. A profile ID is persisted and must
|
||||
/// never carry an account identity.
|
||||
const fn profile_id(self) -> &'static str {
|
||||
match self {
|
||||
Self::Primary => "codex-app-server-primary",
|
||||
Self::Secondary => "codex-app-server-secondary",
|
||||
}
|
||||
}
|
||||
|
||||
/// Codex rate limits are account-scoped, not model-scoped. Naming the
|
||||
/// scope is honest; inventing a model name would not be.
|
||||
const fn scope(self) -> &'static str {
|
||||
match self {
|
||||
Self::Primary => "primary window",
|
||||
Self::Secondary => "secondary window",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the session wants the caller to do next.
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum SessionStep {
|
||||
/// Nothing to do.
|
||||
Idle,
|
||||
/// Write this line to the child's stdin.
|
||||
Send(String),
|
||||
/// Record these observations.
|
||||
Observations(Vec<UsageObservation>),
|
||||
/// The probe's liveness changed.
|
||||
Health(ProbeHealth),
|
||||
/// A server-to-client request was refused. Write the line, and note why.
|
||||
Refused {
|
||||
line: String,
|
||||
class: ServerRequestClass,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
enum Stage {
|
||||
AwaitingInitialize,
|
||||
Ready,
|
||||
}
|
||||
|
||||
/// The last thing recorded for one profile, used to suppress identical rows.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
struct LastRecorded {
|
||||
permille: u64,
|
||||
window: Option<UsageWindow>,
|
||||
at_ms: u64,
|
||||
}
|
||||
|
||||
pub(crate) struct CodexSession {
|
||||
profiles: Vec<AccountProfile>,
|
||||
stage: Stage,
|
||||
next_id: u64,
|
||||
initialize_id: u64,
|
||||
pending_read_id: Option<u64>,
|
||||
last: [Option<LastRecorded>; 2],
|
||||
refused_credential_requests: u64,
|
||||
}
|
||||
|
||||
impl CodexSession {
|
||||
pub(crate) fn new() -> Result<Self, HarnessError> {
|
||||
let profiles = WindowSlot::ALL
|
||||
.into_iter()
|
||||
.map(|slot| {
|
||||
AccountProfile::new(
|
||||
slot.profile_id(),
|
||||
"Codex",
|
||||
"ChatGPT",
|
||||
slot.scope(),
|
||||
"subscription",
|
||||
)
|
||||
.map_err(|_| HarnessError::EmptyProgram)
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(Self {
|
||||
profiles,
|
||||
stage: Stage::AwaitingInitialize,
|
||||
next_id: 1,
|
||||
initialize_id: 0,
|
||||
pending_read_id: None,
|
||||
last: [None, None],
|
||||
refused_credential_requests: 0,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn profiles(&self) -> &[AccountProfile] {
|
||||
&self.profiles
|
||||
}
|
||||
|
||||
/// Whether the handshake has completed and readings may be requested.
|
||||
pub(crate) const fn is_ready(&self) -> bool {
|
||||
matches!(self.stage, Stage::Ready)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) const fn refused_credential_requests(&self) -> u64 {
|
||||
self.refused_credential_requests
|
||||
}
|
||||
|
||||
fn take_id(&mut self) -> u64 {
|
||||
let id = self.next_id;
|
||||
self.next_id = self.next_id.saturating_add(1);
|
||||
id
|
||||
}
|
||||
|
||||
/// The handshake frame. Codex requires `initialize` before any other call.
|
||||
pub(crate) fn initialize(&mut self, client_name: &str, client_version: &str) -> String {
|
||||
let id = self.take_id();
|
||||
self.initialize_id = id;
|
||||
encode_request(
|
||||
id,
|
||||
OutboundRequest::Initialize,
|
||||
&json!({ "clientInfo": { "name": client_name, "version": client_version } }),
|
||||
)
|
||||
}
|
||||
|
||||
/// Requests a fresh reading, unless one is already outstanding.
|
||||
pub(crate) fn poll_frame(&mut self) -> Option<String> {
|
||||
if self.stage != Stage::Ready || self.pending_read_id.is_some() {
|
||||
return None;
|
||||
}
|
||||
let id = self.take_id();
|
||||
self.pending_read_id = Some(id);
|
||||
Some(encode_request(
|
||||
id,
|
||||
OutboundRequest::ReadAccountRateLimits,
|
||||
&json!({}),
|
||||
))
|
||||
}
|
||||
|
||||
/// Consumes one line from the child.
|
||||
pub(crate) fn on_line(&mut self, line: &str, observed_at_ms: u64) -> SessionStep {
|
||||
let Some(inbound) = classify(line) else {
|
||||
return SessionStep::Health(ProbeHealth::Faulted(HarnessError::Malformed));
|
||||
};
|
||||
match inbound {
|
||||
Inbound::ServerRequest { id, class } => {
|
||||
if class == ServerRequestClass::CredentialRefresh {
|
||||
self.refused_credential_requests =
|
||||
self.refused_credential_requests.saturating_add(1);
|
||||
}
|
||||
SessionStep::Refused {
|
||||
line: encode_refusal(&id),
|
||||
class,
|
||||
}
|
||||
}
|
||||
Inbound::Result { id, result } => self.on_result(id, &result, observed_at_ms),
|
||||
Inbound::Failure { id, code } => self.on_failure(id, code, observed_at_ms),
|
||||
// Matching the kind exhaustively means adding a second accepted
|
||||
// notification cannot silently reuse the rate-limit path.
|
||||
Inbound::Notification {
|
||||
kind: AcceptedNotification::AccountRateLimitsUpdated,
|
||||
params,
|
||||
} => self.on_snapshot_body(¶ms, observed_at_ms, true),
|
||||
Inbound::Ignored => SessionStep::Idle,
|
||||
}
|
||||
}
|
||||
|
||||
fn on_result(&mut self, id: u64, result: &Value, observed_at_ms: u64) -> SessionStep {
|
||||
if self.stage == Stage::AwaitingInitialize && id == self.initialize_id {
|
||||
self.stage = Stage::Ready;
|
||||
return SessionStep::Send(encode_notification("initialized", &json!({})));
|
||||
}
|
||||
if self.pending_read_id == Some(id) {
|
||||
self.pending_read_id = None;
|
||||
return self.on_snapshot_body(result, observed_at_ms, false);
|
||||
}
|
||||
SessionStep::Idle
|
||||
}
|
||||
|
||||
fn on_failure(&mut self, id: u64, code: i64, observed_at_ms: u64) -> SessionStep {
|
||||
if self.pending_read_id == Some(id) {
|
||||
self.pending_read_id = None;
|
||||
// Codex answers -32600 when the account cannot report a quota,
|
||||
// which is a true statement about the account rather than a fault.
|
||||
if code == INVALID_REQUEST {
|
||||
let observations = self.unavailable_for_all(observed_at_ms);
|
||||
self.forget_all();
|
||||
return SessionStep::Observations(observations);
|
||||
}
|
||||
return SessionStep::Health(ProbeHealth::Faulted(HarnessError::Rejected(code)));
|
||||
}
|
||||
if id == self.initialize_id {
|
||||
return SessionStep::Health(ProbeHealth::Faulted(HarnessError::Rejected(code)));
|
||||
}
|
||||
SessionStep::Idle
|
||||
}
|
||||
|
||||
/// Turns a `rateLimits` value into observations.
|
||||
///
|
||||
/// `sparse` marks a push notification. Upstream documents that a rolling
|
||||
/// update may omit values and that omission "does not clear a previously
|
||||
/// observed value", so an absent window on a push records nothing at all.
|
||||
/// It must not record an explicit gap, and it must not be merged into the
|
||||
/// last full read either: a merged composite would look provider-reported
|
||||
/// while being partly a memory.
|
||||
fn on_snapshot_body(&mut self, body: &Value, observed_at_ms: u64, sparse: bool) -> SessionStep {
|
||||
let Ok(body) = serde_json::from_value::<RateLimitsResultWire>(body.clone()) else {
|
||||
return SessionStep::Health(ProbeHealth::Faulted(HarnessError::Malformed));
|
||||
};
|
||||
let observations = self.observations_from(&body.rate_limits, observed_at_ms, sparse);
|
||||
if observations.is_empty() {
|
||||
SessionStep::Idle
|
||||
} else {
|
||||
SessionStep::Observations(observations)
|
||||
}
|
||||
}
|
||||
|
||||
fn observations_from(
|
||||
&mut self,
|
||||
snapshot: &RateLimitSnapshotWire,
|
||||
observed_at_ms: u64,
|
||||
sparse: bool,
|
||||
) -> Vec<UsageObservation> {
|
||||
let mut observations = Vec::new();
|
||||
for slot in WindowSlot::ALL {
|
||||
let index = slot as usize;
|
||||
let wire = match slot {
|
||||
WindowSlot::Primary => snapshot.primary.as_ref(),
|
||||
WindowSlot::Secondary => snapshot.secondary.as_ref(),
|
||||
};
|
||||
let Some(wire) = wire else {
|
||||
if !sparse {
|
||||
// A full read that omits a window means the window is gone.
|
||||
if self.last[index].take().is_some() {
|
||||
observations.push(UsageObservation::unavailable(
|
||||
self.profile_id(index),
|
||||
observed_at_ms,
|
||||
));
|
||||
}
|
||||
}
|
||||
continue;
|
||||
};
|
||||
let reading = parse_window(wire, observed_at_ms);
|
||||
if !self.should_record(index, reading, observed_at_ms) {
|
||||
continue;
|
||||
}
|
||||
observations.push(to_observation(
|
||||
&self.profile_id(index),
|
||||
reading,
|
||||
observed_at_ms,
|
||||
));
|
||||
}
|
||||
observations
|
||||
}
|
||||
|
||||
fn profile_id(&self, index: usize) -> AccountProfileId {
|
||||
self.profiles[index].id().clone()
|
||||
}
|
||||
|
||||
/// Suppresses an identical reading unless the re-emit floor has passed.
|
||||
fn should_record(&mut self, index: usize, reading: WindowReading, observed_at_ms: u64) -> bool {
|
||||
let WindowReading::Usable { permille, window } = reading else {
|
||||
// A gap is only worth recording when it changes the story.
|
||||
return self.last[index].take().is_some();
|
||||
};
|
||||
let unchanged = self.last[index].is_some_and(|last| {
|
||||
last.permille == permille
|
||||
&& last.window == window
|
||||
&& observed_at_ms.saturating_sub(last.at_ms) < MIN_REEMIT_INTERVAL_MS
|
||||
});
|
||||
if unchanged {
|
||||
return false;
|
||||
}
|
||||
self.last[index] = Some(LastRecorded {
|
||||
permille,
|
||||
window,
|
||||
at_ms: observed_at_ms,
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
/// An explicit gap for every profile, stamped with when the gap was
|
||||
/// observed.
|
||||
///
|
||||
/// Stamping these from the last successful reading — or from the epoch when
|
||||
/// there has never been one — would put them behind the ledger's newest
|
||||
/// entry, and an append-only stream rejects a backwards timestamp. The gap
|
||||
/// would then be silently dropped and a stale percentage would keep
|
||||
/// rendering as current, which is the exact failure decision 0012 exists to
|
||||
/// prevent.
|
||||
fn unavailable_for_all(&self, observed_at_ms: u64) -> Vec<UsageObservation> {
|
||||
(0..self.profiles.len())
|
||||
.map(|index| UsageObservation::unavailable(self.profile_id(index), observed_at_ms))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn forget_all(&mut self) {
|
||||
self.last = [None, None];
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{CodexSession, SessionStep, WindowSlot};
|
||||
use crate::jsonrpc::ServerRequestClass;
|
||||
use crate::probe::ProbeHealth;
|
||||
use lumbridge_core::{UsageProvenance, UsageUnit};
|
||||
|
||||
const HOUR_MS: u64 = 3_600_000;
|
||||
const RESETS_AT_SECONDS: i64 = 1_730_947_200;
|
||||
const RESETS_AT_MS: u64 = 1_730_947_200_000;
|
||||
|
||||
fn session() -> CodexSession {
|
||||
CodexSession::new().expect("the session profiles are valid")
|
||||
}
|
||||
|
||||
fn ready() -> CodexSession {
|
||||
let mut session = session();
|
||||
let _ = session.initialize("lumbridge", "0.0.1");
|
||||
let step = session.on_line(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#, 0);
|
||||
assert!(matches!(step, SessionStep::Send(_)), "handshake completes");
|
||||
session
|
||||
}
|
||||
|
||||
fn read_response(id: u64, percent: i64, minutes: i64) -> String {
|
||||
format!(
|
||||
r#"{{"jsonrpc":"2.0","id":{id},"result":{{"rateLimits":{{"primary":{{"usedPercent":{percent},"windowDurationMins":{minutes},"resetsAt":{RESETS_AT_SECONDS}}},"secondary":null}}}}}}"#
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_two_windows_are_two_separate_account_free_profiles() {
|
||||
let session = session();
|
||||
let ids = session
|
||||
.profiles()
|
||||
.iter()
|
||||
.map(|profile| profile.id().as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
ids,
|
||||
["codex-app-server-primary", "codex-app-server-secondary"]
|
||||
);
|
||||
assert!(
|
||||
!ids.iter().any(|id| id.contains('@')),
|
||||
"a persisted profile ID must not carry an account identity"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_reading_is_requested_before_the_handshake_completes() {
|
||||
let mut session = session();
|
||||
let _ = session.initialize("lumbridge", "0.0.1");
|
||||
assert!(
|
||||
session.poll_frame().is_none(),
|
||||
"the app-server requires initialize first"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn one_read_is_outstanding_at_a_time() {
|
||||
let mut session = ready();
|
||||
assert!(session.poll_frame().is_some());
|
||||
assert!(
|
||||
session.poll_frame().is_none(),
|
||||
"a second read must wait for the first to answer"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_full_read_becomes_a_provider_reported_observation() {
|
||||
let mut session = ready();
|
||||
let frame = session.poll_frame().expect("ready sessions poll");
|
||||
assert!(frame.contains("account/rateLimits/read"));
|
||||
let step = session.on_line(&read_response(2, 25, 300), RESETS_AT_MS - HOUR_MS);
|
||||
let SessionStep::Observations(observations) = step else {
|
||||
panic!("a full read produces observations");
|
||||
};
|
||||
assert_eq!(observations.len(), 1, "only primary was present");
|
||||
let observation = &observations[0];
|
||||
assert_eq!(observation.provenance(), UsageProvenance::ProviderReported);
|
||||
assert_eq!(observation.unit(), Some(UsageUnit::WindowPermille));
|
||||
assert_eq!(observation.consumed(), 250);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unchanged_reading_is_not_re_recorded_every_poll() {
|
||||
let mut session = ready();
|
||||
let _ = session.poll_frame();
|
||||
let at = RESETS_AT_MS - HOUR_MS;
|
||||
assert!(matches!(
|
||||
session.on_line(&read_response(2, 25, 300), at),
|
||||
SessionStep::Observations(_)
|
||||
));
|
||||
let _ = session.poll_frame();
|
||||
assert!(
|
||||
matches!(
|
||||
session.on_line(&read_response(3, 25, 300), at + 60_000),
|
||||
SessionStep::Idle
|
||||
),
|
||||
"an identical reading must not flush the bounded retention"
|
||||
);
|
||||
let _ = session.poll_frame();
|
||||
assert!(
|
||||
matches!(
|
||||
session.on_line(&read_response(4, 26, 300), at + 120_000),
|
||||
SessionStep::Observations(_)
|
||||
),
|
||||
"a changed reading is always recorded"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unchanged_reading_is_re_emitted_after_the_floor_so_a_rate_stays_derivable() {
|
||||
let mut session = ready();
|
||||
let _ = session.poll_frame();
|
||||
let at = RESETS_AT_MS - 4 * HOUR_MS;
|
||||
let _ = session.on_line(&read_response(2, 25, 300), at);
|
||||
let _ = session.poll_frame();
|
||||
assert!(matches!(
|
||||
session.on_line(&read_response(3, 25, 300), at + 300_001),
|
||||
SessionStep::Observations(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_sparse_push_without_a_window_records_nothing() {
|
||||
let mut session = ready();
|
||||
let _ = session.poll_frame();
|
||||
let at = RESETS_AT_MS - HOUR_MS;
|
||||
let _ = session.on_line(&read_response(2, 25, 300), at);
|
||||
// Upstream: nullable metadata absent from a rolling update "does not
|
||||
// clear a previously observed value".
|
||||
let push =
|
||||
r#"{"jsonrpc":"2.0","method":"account/rateLimits/updated","params":{"rateLimits":{}}}"#;
|
||||
assert!(
|
||||
matches!(session.on_line(push, at + 1_000), SessionStep::Idle),
|
||||
"a sparse push must neither clear nor merge"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_full_read_that_drops_a_window_records_an_explicit_gap() {
|
||||
let mut session = ready();
|
||||
let _ = session.poll_frame();
|
||||
let at = RESETS_AT_MS - HOUR_MS;
|
||||
let _ = session.on_line(&read_response(2, 25, 300), at);
|
||||
let _ = session.poll_frame();
|
||||
let empty =
|
||||
r#"{"jsonrpc":"2.0","id":3,"result":{"rateLimits":{"primary":null,"secondary":null}}}"#;
|
||||
let SessionStep::Observations(observations) = session.on_line(empty, at + 1_000) else {
|
||||
panic!("a full read that loses a window is a change worth recording");
|
||||
};
|
||||
assert_eq!(observations.len(), 1);
|
||||
assert_eq!(observations[0].provenance(), UsageProvenance::Unavailable);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_denial_stamps_its_gap_with_now_so_the_ledger_accepts_it() {
|
||||
use lumbridge_core::UsageLedger;
|
||||
|
||||
let mut session = ready();
|
||||
let _ = session.poll_frame();
|
||||
let at = RESETS_AT_MS - HOUR_MS;
|
||||
let SessionStep::Observations(good) = session.on_line(&read_response(2, 25, 300), at)
|
||||
else {
|
||||
panic!("the first read succeeds");
|
||||
};
|
||||
let _ = session.poll_frame();
|
||||
let denied = r#"{"jsonrpc":"2.0","id":3,"error":{"code":-32600,"message":"chatgpt authentication required to read rate limits"}}"#;
|
||||
let SessionStep::Observations(gaps) = session.on_line(denied, at + 60_000) else {
|
||||
panic!("a denial is a fact about the account");
|
||||
};
|
||||
|
||||
// The ledger is append-only, so a gap stamped behind the reading it
|
||||
// supersedes would be refused and the stale reading would survive.
|
||||
let mut ledger = UsageLedger::new();
|
||||
for observation in good.into_iter().chain(gaps) {
|
||||
ledger
|
||||
.record(observation)
|
||||
.expect("every observation must be accepted in order");
|
||||
}
|
||||
let projection = ledger.project(session.profiles()[0].id(), at + 60_000);
|
||||
assert!(
|
||||
!projection.is_available(),
|
||||
"the denial must supersede the earlier reading"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_api_key_account_reports_unavailable_rather_than_a_fault() {
|
||||
let mut session = ready();
|
||||
let _ = session.poll_frame();
|
||||
let denied = r#"{"jsonrpc":"2.0","id":2,"error":{"code":-32600,"message":"chatgpt authentication required to read rate limits"}}"#;
|
||||
let SessionStep::Observations(observations) = session.on_line(denied, 5_000) else {
|
||||
panic!("a quota-less account is a fact about the account");
|
||||
};
|
||||
assert_eq!(observations.len(), 2);
|
||||
assert!(
|
||||
observations
|
||||
.iter()
|
||||
.all(|observation| observation.provenance() == UsageProvenance::Unavailable)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn any_other_rejection_faults_the_probe() {
|
||||
let mut session = ready();
|
||||
let _ = session.poll_frame();
|
||||
let denied = r#"{"jsonrpc":"2.0","id":2,"error":{"code":-32603,"message":"boom"}}"#;
|
||||
assert!(matches!(
|
||||
session.on_line(denied, 5_000),
|
||||
SessionStep::Health(ProbeHealth::Faulted(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_credential_refresh_request_is_refused_and_counted() {
|
||||
let mut session = ready();
|
||||
let step = session.on_line(
|
||||
r#"{"jsonrpc":"2.0","id":99,"method":"account/chatgptAuthTokens/refresh"}"#,
|
||||
0,
|
||||
);
|
||||
let SessionStep::Refused { line, class } = step else {
|
||||
panic!("a credential request must be refused");
|
||||
};
|
||||
assert_eq!(class, ServerRequestClass::CredentialRefresh);
|
||||
assert!(line.contains("-32601"));
|
||||
assert!(
|
||||
!line.contains("access"),
|
||||
"a refusal must not echo the request"
|
||||
);
|
||||
assert_eq!(session.refused_credential_requests(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_malformed_line_faults_rather_than_being_guessed_at() {
|
||||
let mut session = ready();
|
||||
assert!(matches!(
|
||||
session.on_line("this is not json", 0),
|
||||
SessionStep::Health(ProbeHealth::Faulted(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn window_slots_have_distinct_static_identifiers() {
|
||||
assert_ne!(
|
||||
WindowSlot::Primary.profile_id(),
|
||||
WindowSlot::Secondary.profile_id()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
//! A deliberately deaf JSON-RPC client.
|
||||
//!
|
||||
//! The Codex app-server can send requests *to* its client, including
|
||||
//! `account/chatgptAuthTokens/refresh`, which returns a bare access token.
|
||||
//! AGENTS.md forbids Lumbridge from scraping a harness's private credentials.
|
||||
//! Choosing not to call that method would be a convention; this module makes
|
||||
//! it unrepresentable instead.
|
||||
//!
|
||||
//! There is no free-form method string anywhere in this crate. Outbound
|
||||
//! requests come from a two-variant enum, and every inbound server request is
|
||||
//! answered with `-32601 method not found` regardless of what it asks for.
|
||||
//! Classification exists only so a refusal can be counted and named, which
|
||||
//! makes "we were asked for a credential and refused" an auditable event
|
||||
//! rather than an absence.
|
||||
|
||||
use serde_json::{Value, json};
|
||||
|
||||
/// JSON-RPC's "method not found". The only reply this client ever sends.
|
||||
const METHOD_NOT_FOUND: i64 = -32601;
|
||||
/// Codex returns this for a request that needs an account it does not have.
|
||||
pub(crate) const INVALID_REQUEST: i64 = -32600;
|
||||
|
||||
/// Every request this crate is able to send.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum OutboundRequest {
|
||||
Initialize,
|
||||
ReadAccountRateLimits,
|
||||
}
|
||||
|
||||
impl OutboundRequest {
|
||||
pub(crate) const fn method(self) -> &'static str {
|
||||
match self {
|
||||
Self::Initialize => "initialize",
|
||||
Self::ReadAccountRateLimits => "account/rateLimits/read",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The notifications this client accepts. Anything else is ignored.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum AcceptedNotification {
|
||||
AccountRateLimitsUpdated,
|
||||
}
|
||||
|
||||
impl AcceptedNotification {
|
||||
pub(crate) fn from_method(method: &str) -> Option<Self> {
|
||||
match method {
|
||||
"account/rateLimits/updated" => Some(Self::AccountRateLimitsUpdated),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What a server-to-client request was asking for.
|
||||
///
|
||||
/// Purely descriptive. Every class is refused identically; the distinction
|
||||
/// exists so a probe can report that it declined a credential request.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum ServerRequestClass {
|
||||
/// A request that would hand Lumbridge a credential. Always refused.
|
||||
CredentialRefresh,
|
||||
/// A request to run something on the harness's behalf.
|
||||
Execution,
|
||||
/// A request for a human approval decision.
|
||||
Approval,
|
||||
/// A request for ambient information such as the current time.
|
||||
Ambient,
|
||||
/// Anything this client does not recognise.
|
||||
Unrecognised,
|
||||
}
|
||||
|
||||
impl ServerRequestClass {
|
||||
pub(crate) fn classify(method: &str) -> Self {
|
||||
match method {
|
||||
"account/chatgptAuthTokens/refresh" => Self::CredentialRefresh,
|
||||
"item/tool/call" | "attestation/generate" => Self::Execution,
|
||||
"item/permissions/requestApproval" => Self::Approval,
|
||||
"currentTime/read" => Self::Ambient,
|
||||
_ => Self::Unrecognised,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn label(self) -> &'static str {
|
||||
match self {
|
||||
Self::CredentialRefresh => "credential refresh",
|
||||
Self::Execution => "execution",
|
||||
Self::Approval => "approval",
|
||||
Self::Ambient => "ambient information",
|
||||
Self::Unrecognised => "unrecognised request",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A classified inbound line.
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum Inbound {
|
||||
/// A successful response to one of our two requests.
|
||||
Result { id: u64, result: Value },
|
||||
/// A failed response to one of our two requests.
|
||||
Failure { id: u64, code: i64 },
|
||||
/// An accepted push notification.
|
||||
Notification {
|
||||
kind: AcceptedNotification,
|
||||
params: Value,
|
||||
},
|
||||
/// The server asked us for something. It will be refused.
|
||||
ServerRequest {
|
||||
id: Value,
|
||||
class: ServerRequestClass,
|
||||
},
|
||||
/// A well-formed frame with nothing for us in it.
|
||||
Ignored,
|
||||
}
|
||||
|
||||
/// Encodes one of the two permitted requests.
|
||||
pub(crate) fn encode_request(id: u64, request: OutboundRequest, params: &Value) -> String {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"method": request.method(),
|
||||
"params": params,
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
pub(crate) fn encode_notification(method: &'static str, params: &Value) -> String {
|
||||
json!({ "jsonrpc": "2.0", "method": method, "params": params }).to_string()
|
||||
}
|
||||
|
||||
/// The refusal sent for every server-to-client request.
|
||||
pub(crate) fn encode_refusal(id: &Value) -> String {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"error": { "code": METHOD_NOT_FOUND, "message": "method not found" },
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Classifies one line of JSON-RPC without interpreting its payload.
|
||||
///
|
||||
/// Returns `None` when the line is not a JSON object at all, which the caller
|
||||
/// reports as [`crate::HarnessError::Malformed`].
|
||||
pub(crate) fn classify(line: &str) -> Option<Inbound> {
|
||||
let frame: Value = serde_json::from_str(line).ok()?;
|
||||
let object = frame.as_object()?;
|
||||
let method = object.get("method").and_then(Value::as_str);
|
||||
let id = object.get("id");
|
||||
|
||||
match (method, id) {
|
||||
// A method with an id is the server asking us for something.
|
||||
(Some(method), Some(id)) => Some(Inbound::ServerRequest {
|
||||
id: id.clone(),
|
||||
class: ServerRequestClass::classify(method),
|
||||
}),
|
||||
// A method without an id is a notification.
|
||||
(Some(method), None) => Some(AcceptedNotification::from_method(method).map_or(
|
||||
Inbound::Ignored,
|
||||
|kind| Inbound::Notification {
|
||||
kind,
|
||||
params: object.get("params").cloned().unwrap_or(Value::Null),
|
||||
},
|
||||
)),
|
||||
// An id without a method is a response to one of ours.
|
||||
(None, Some(id)) => {
|
||||
let id = id.as_u64()?;
|
||||
if let Some(error) = object.get("error") {
|
||||
let code = error.get("code").and_then(Value::as_i64).unwrap_or(0);
|
||||
return Some(Inbound::Failure { id, code });
|
||||
}
|
||||
Some(Inbound::Result {
|
||||
id,
|
||||
result: object.get("result").cloned().unwrap_or(Value::Null),
|
||||
})
|
||||
}
|
||||
(None, None) => Some(Inbound::Ignored),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
AcceptedNotification, Inbound, OutboundRequest, ServerRequestClass, classify,
|
||||
encode_refusal, encode_request,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
#[test]
|
||||
fn only_two_methods_can_be_sent() {
|
||||
assert_eq!(OutboundRequest::Initialize.method(), "initialize");
|
||||
assert_eq!(
|
||||
OutboundRequest::ReadAccountRateLimits.method(),
|
||||
"account/rateLimits/read"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_credential_refresh_request_is_classified_and_refused() {
|
||||
let line = r#"{"jsonrpc":"2.0","id":9,"method":"account/chatgptAuthTokens/refresh"}"#;
|
||||
let Some(Inbound::ServerRequest { id, class }) = classify(line) else {
|
||||
panic!("a method with an id is a server request");
|
||||
};
|
||||
assert_eq!(class, ServerRequestClass::CredentialRefresh);
|
||||
let refusal: Value =
|
||||
serde_json::from_str(&encode_refusal(&id)).expect("the refusal is valid JSON");
|
||||
assert_eq!(refusal["error"]["code"], json!(-32601));
|
||||
assert!(
|
||||
refusal.get("result").is_none(),
|
||||
"a refusal must never carry a result"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_server_request_class_is_refused_the_same_way() {
|
||||
for method in [
|
||||
"account/chatgptAuthTokens/refresh",
|
||||
"item/tool/call",
|
||||
"item/permissions/requestApproval",
|
||||
"currentTime/read",
|
||||
"something/entirely/new",
|
||||
] {
|
||||
let line = format!(r#"{{"jsonrpc":"2.0","id":1,"method":"{method}"}}"#);
|
||||
let Some(Inbound::ServerRequest { id, .. }) = classify(&line) else {
|
||||
panic!("{method} must classify as a server request");
|
||||
};
|
||||
let refusal: Value = serde_json::from_str(&encode_refusal(&id)).expect("valid JSON");
|
||||
assert_eq!(refusal["error"]["code"], json!(-32601));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_the_rate_limit_notification_is_accepted() {
|
||||
let accepted = classify(
|
||||
r#"{"jsonrpc":"2.0","method":"account/rateLimits/updated","params":{"rateLimits":{}}}"#,
|
||||
);
|
||||
assert!(matches!(
|
||||
accepted,
|
||||
Some(Inbound::Notification {
|
||||
kind: AcceptedNotification::AccountRateLimitsUpdated,
|
||||
..
|
||||
})
|
||||
));
|
||||
let ignored =
|
||||
classify(r#"{"jsonrpc":"2.0","method":"thread/tokenUsage/updated","params":{}}"#);
|
||||
assert!(matches!(ignored, Some(Inbound::Ignored)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_and_failures_are_separated() {
|
||||
assert!(matches!(
|
||||
classify(r#"{"jsonrpc":"2.0","id":7,"result":{"rateLimits":{}}}"#),
|
||||
Some(Inbound::Result { id: 7, .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
classify(r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32600,"message":"nope"}}"#),
|
||||
Some(Inbound::Failure {
|
||||
id: 7,
|
||||
code: -32600
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_and_non_object_lines_are_rejected() {
|
||||
for line in ["", "not json", "[1,2,3]", "\"a string\"", "{"] {
|
||||
assert!(classify(line).is_none(), "{line:?} must not classify");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_request_encodes_without_a_free_form_method() {
|
||||
let encoded = encode_request(1, OutboundRequest::ReadAccountRateLimits, &json!({}));
|
||||
let frame: Value = serde_json::from_str(&encoded).expect("valid JSON");
|
||||
assert_eq!(frame["method"], json!("account/rateLimits/read"));
|
||||
assert_eq!(frame["jsonrpc"], json!("2.0"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
//! Documented status and usage probes for supervised coding harnesses.
|
||||
//!
|
||||
//! `lumbridge-core` owns the usage ledger and stays free of IO, clocks, and
|
||||
//! async. This crate is the other side of that boundary: it runs processes,
|
||||
//! reads a clock, parses untrusted wire text, and hands back
|
||||
//! [`lumbridge_core::UsageObservation`] values. Nothing here interprets a
|
||||
//! usage number; it only decides which observation is honest to record.
|
||||
//!
|
||||
//! The rule every probe follows: Lumbridge observes documented surfaces of a
|
||||
//! harness it launched, and never reads that harness's credentials. A probe
|
||||
//! that cannot obtain a reading records
|
||||
//! [`lumbridge_core::UsageObservation::unavailable`] rather than a zero.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
mod child;
|
||||
pub mod claude;
|
||||
mod clock;
|
||||
pub mod codex;
|
||||
mod jsonrpc;
|
||||
mod probe;
|
||||
|
||||
pub use child::PROBE_ENV_ALLOWLIST;
|
||||
pub use clock::MonotonicWallClock;
|
||||
pub use jsonrpc::ServerRequestClass;
|
||||
pub use probe::{ProbeHealth, ProbeOutcome, UsageProbe};
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// Failures a probe can report.
|
||||
///
|
||||
/// Every variant is [`Copy`] and carries only a static discriminator, an
|
||||
/// [`std::io::ErrorKind`], or a number. No variant can capture a string from a
|
||||
/// child process or the filesystem, so an error can never smuggle a path,
|
||||
/// a credential, or attacker-chosen text into a log or a UI surface.
|
||||
#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
|
||||
pub enum HarnessError {
|
||||
#[error("a probe program name must not be empty")]
|
||||
EmptyProgram,
|
||||
#[error("the environment variable {0} is not on the probe allowlist")]
|
||||
EnvNotAllowlisted(&'static str),
|
||||
#[error("a probe poll interval must be at least one second")]
|
||||
ProbeIntervalTooShort,
|
||||
#[error("the probe program could not be started ({0:?})")]
|
||||
SpawnFailed(std::io::ErrorKind),
|
||||
#[error("the probe program closed its output")]
|
||||
Eof,
|
||||
#[error("the probe program sent a {0} byte line, over the accepted limit")]
|
||||
OversizedLine(usize),
|
||||
#[error("the probe program sent a line that is not a JSON-RPC frame")]
|
||||
Malformed,
|
||||
#[error("the harness rejected the request with JSON-RPC code {0}")]
|
||||
Rejected(i64),
|
||||
#[error("the harness has no subscription account to report a quota for")]
|
||||
NotAuthenticated,
|
||||
}
|
||||
|
||||
impl HarnessError {
|
||||
/// Whether this failure means "there is no number", as opposed to
|
||||
/// "the probe broke".
|
||||
///
|
||||
/// Both record an unavailable observation, but only the second is a fault
|
||||
/// worth surfacing as a broken probe.
|
||||
#[must_use]
|
||||
pub const fn is_expected_gap(self) -> bool {
|
||||
matches!(self, Self::NotAuthenticated)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
//! The contract between a harness probe and the usage ledger.
|
||||
//!
|
||||
//! This trait lives here rather than in `lumbridge-core` because implementing
|
||||
//! it implies a fallible external call, and core stays IO-free, clock-free,
|
||||
//! and async-free. [`lumbridge_core::UsageObservation`] is already the shared
|
||||
//! type; nothing further needs to cross the boundary.
|
||||
|
||||
use lumbridge_core::{AccountProfile, UsageObservation};
|
||||
|
||||
use crate::HarnessError;
|
||||
|
||||
/// What a probe knows about its own liveness.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum ProbeHealth {
|
||||
/// Launched, handshake not finished.
|
||||
Starting,
|
||||
/// Handshake complete and readings are expected.
|
||||
Ready,
|
||||
/// Running, but this account has no subscription quota to report. This is
|
||||
/// a correct terminal state, not a fault: API-key users have no window.
|
||||
NoQuotaAccount,
|
||||
/// Broken. The observations it produced before this are still facts.
|
||||
Faulted(HarnessError),
|
||||
/// Deliberately shut down.
|
||||
Stopped,
|
||||
}
|
||||
|
||||
impl ProbeHealth {
|
||||
#[must_use]
|
||||
pub const fn is_faulted(self) -> bool {
|
||||
matches!(self, Self::Faulted(_))
|
||||
}
|
||||
|
||||
/// A short phrase a UI may render beside the probe's profiles.
|
||||
#[must_use]
|
||||
pub const fn label(self) -> &'static str {
|
||||
match self {
|
||||
Self::Starting => "probe starting",
|
||||
Self::Ready => "probe ready",
|
||||
Self::NoQuotaAccount => "no subscription quota on this account",
|
||||
Self::Faulted(_) => "probe faulted",
|
||||
Self::Stopped => "probe stopped",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One round of probe output.
|
||||
///
|
||||
/// It carries observations and liveness and nothing else. There is no variant
|
||||
/// able to carry a command, a plan, or a capability request, so a probe cannot
|
||||
/// become an actor on the workspace no matter what the harness sends it.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ProbeOutcome {
|
||||
observations: Vec<UsageObservation>,
|
||||
health: ProbeHealth,
|
||||
}
|
||||
|
||||
impl ProbeOutcome {
|
||||
#[must_use]
|
||||
pub const fn new(observations: Vec<UsageObservation>, health: ProbeHealth) -> Self {
|
||||
Self {
|
||||
observations,
|
||||
health,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn idle(health: ProbeHealth) -> Self {
|
||||
Self {
|
||||
observations: Vec::new(),
|
||||
health,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn observations(&self) -> &[UsageObservation] {
|
||||
&self.observations
|
||||
}
|
||||
|
||||
/// Takes the observations out for recording, leaving the outcome empty.
|
||||
#[must_use]
|
||||
pub fn into_observations(self) -> Vec<UsageObservation> {
|
||||
self.observations
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn health(&self) -> ProbeHealth {
|
||||
self.health
|
||||
}
|
||||
}
|
||||
|
||||
/// A source of usage observations for one or more account profiles.
|
||||
pub trait UsageProbe {
|
||||
/// Drains whatever the probe has ready.
|
||||
///
|
||||
/// Implementations must return promptly: this is called from the UI's
|
||||
/// frame loop and must never wait on a child process, a socket, or a lock
|
||||
/// held across IO.
|
||||
fn poll(&mut self) -> ProbeOutcome;
|
||||
|
||||
/// The profiles this probe can produce observations for.
|
||||
fn profiles(&self) -> &[AccountProfile];
|
||||
|
||||
/// Stops the probe and releases its child process.
|
||||
fn shutdown(&mut self);
|
||||
}
|
||||
Reference in New Issue
Block a user