lumbridge-code/crates/lumbridge-harness/src/codex/session.rs

//! 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(&params, 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()
        );
    }
}