af-agent-session 0.7.0

//! Append-only corrections to measurement facts. Reservation settlement remains immutable.
use af_context::{MeteringCorrectionId, RunId, SubjectId};
use serde::{Deserialize, Serialize};

use crate::{EventError, MeteringDetails, MeteringSource, SessionProjection};

/// Latest measurement for one stable Run operation; revision zero is its original fact.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OperationUsage {
    /// Prompt tokens measured.
    pub prompt_tokens: u64,
    /// Completion tokens measured.
    pub completion_tokens: u64,
    /// Provider-neutral resource weight, never a price.
    pub cost_units: u64,
    /// Attribution; absent only for legacy data.
    pub metering: Option<MeteringDetails>,
    /// Number of accepted corrections to this operation.
    pub revision: u64,
}

/// Privileged report that replaces the derived measurement, never its original event.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct MeteringCorrection {
    /// Stable idempotency identity, unique within the Session.
    pub id: MeteringCorrectionId,
    /// Terminal Run whose measurement is corrected.
    pub run_id: RunId,
    /// Original stable operation identity.
    pub operation: String,
    /// Exact measurement revision the reporter observed.
    pub expected_revision: u64,
    /// Corrected prompt tokens.
    pub prompt_tokens: u64,
    /// Corrected completion tokens.
    pub completion_tokens: u64,
    /// Corrected resource units; no monetary meaning in Factory.
    pub cost_units: u64,
    /// Reported attribution, preserving the original operation identity.
    pub metering: MeteringDetails,
    /// Short audit reason, without credentials or provider response bodies.
    pub reason: String,
}

/// Accepted correction with server-derived actor and immutable event position.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AppliedMeteringCorrection {
    /// Caller-supplied report.
    pub correction: MeteringCorrection,
    /// Authenticated reporting subject.
    pub actor_id: SubjectId,
    /// Session event sequence of first acceptance.
    pub seq: u64,
}

impl MeteringCorrection {
    /// Validate bounded, non-overflowing reported measurements before any write.
    pub fn validate(&self) -> Result<(), EventError> {
        self.metering.validate()?;
        let valid_text = |value: &str| {
            !value.trim().is_empty() && value.len() <= 512 && !value.chars().any(char::is_control)
        };
        if !valid_text(self.id.as_str())
            || !valid_text(&self.operation)
            || !valid_text(&self.reason)
            || self.expected_revision == u64::MAX
            || self
                .prompt_tokens
                .checked_add(self.completion_tokens)
                .and_then(|v| v.checked_add(self.cost_units))
                .is_none_or(|v| v > i64::MAX as u64)
            || !matches!(
                self.metering,
                MeteringDetails::Model {
                    source: MeteringSource::Reported,
                    ..
                } | MeteringDetails::Tool {
                    source: MeteringSource::Reported,
                    ..
                }
            )
        {
            return Err(EventError::InvalidMetering);
        }
        Ok(())
    }
}

impl SessionProjection {
    /// Latest authoritative measurement. Original quota settlement is queried separately.
    pub fn operation_usage(&self, run_id: &RunId, operation: &str) -> Option<OperationUsage> {
        let key = (run_id.clone(), operation.to_owned());
        if let Some(value) = self.corrected_usage.get(&key) {
            return Some(value.clone());
        }
        if let Some((prompt_tokens, completion_tokens, cost_units, metering)) =
            self.usage_operations.get(&key)
        {
            return Some(OperationUsage {
                prompt_tokens: *prompt_tokens,
                completion_tokens: *completion_tokens,
                cost_units: *cost_units,
                metering: metering.clone(),
                revision: 0,
            });
        }
        self.pending_usage_operations.get(&key).map(
            |(prompt_tokens, completion_tokens, cost_units, metering)| OperationUsage {
                prompt_tokens: *prompt_tokens,
                completion_tokens: *completion_tokens,
                cost_units: *cost_units,
                metering: metering.clone(),
                revision: 0,
            },
        )
    }

    /// Original receipt for an idempotently accepted correction.
    pub fn metering_correction(
        &self,
        id: &MeteringCorrectionId,
    ) -> Option<&AppliedMeteringCorrection> {
        self.metering_corrections.get(id)
    }

    pub(super) fn apply_metering_correction(
        &mut self,
        correction: &MeteringCorrection,
        actor_id: &SubjectId,
        seq: u64,
    ) -> Result<(), EventError> {
        correction.validate()?;
        if let Some(applied) = self.metering_corrections.get(&correction.id) {
            return if applied.correction == *correction && applied.actor_id == *actor_id {
                Ok(())
            } else {
                Err(EventError::UsageConflict(correction.operation.clone()))
            };
        }
        let previous = self
            .operation_usage(&correction.run_id, &correction.operation)
            .ok_or_else(|| EventError::UsageConflict(correction.operation.clone()))?;
        if !self
            .run_status
            .get(&correction.run_id)
            .is_some_and(|status| status.is_terminal())
            || previous.revision != correction.expected_revision
            || !previous
                .metering
                .as_ref()
                .is_some_and(|details| correction.metering.completes(details))
        {
            return Err(EventError::UsageConflict(correction.operation.clone()));
        }
        self.corrected_usage.insert(
            (correction.run_id.clone(), correction.operation.clone()),
            OperationUsage {
                prompt_tokens: correction.prompt_tokens,
                completion_tokens: correction.completion_tokens,
                cost_units: correction.cost_units,
                metering: Some(correction.metering.clone()),
                revision: previous.revision + 1,
            },
        );
        self.metering_corrections.insert(
            correction.id.clone(),
            AppliedMeteringCorrection {
                correction: correction.clone(),
                actor_id: actor_id.clone(),
                seq,
            },
        );
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{MeteringOutcome, RunState, RunStatus};

    fn report() -> MeteringCorrection {
        MeteringCorrection {
            id: "correction".parse().unwrap(),
            run_id: "run".parse().unwrap(),
            operation: "operation".into(),
            expected_revision: 0,
            prompt_tokens: 1,
            completion_tokens: 2,
            cost_units: 0,
            metering: MeteringDetails::Model {
                model: "m".into(),
                provider: Some("p".into()),
                provider_attempt_id: "attempt".parse().unwrap(),
                source: MeteringSource::Reported,
                outcome: MeteringOutcome::Completed,
            },
            reason: "provider report".into(),
        }
    }

    #[test]
    fn invalid_reports_and_mismatched_identities_fail_closed() {
        let good = report();
        good.validate().unwrap();
        for mutation in 0..7 {
            let mut invalid = good.clone();
            match mutation {
                0 => invalid.reason.clear(),
                1 => invalid.operation = "a".repeat(513),
                2 => invalid.prompt_tokens = u64::MAX,
                3 => invalid.expected_revision = u64::MAX,
                4 => invalid.id = "bad\nid".parse().unwrap(),
                5 => {
                    if let MeteringDetails::Model { source, .. } = &mut invalid.metering {
                        *source = MeteringSource::Estimated;
                    }
                }
                _ => {
                    if let MeteringDetails::Model { model, .. } = &mut invalid.metering {
                        model.clear();
                    }
                }
            }
            assert!(invalid.validate().is_err());
        }
        let mut projection = SessionProjection::default();
        let actor = "admin".parse().unwrap();
        assert!(projection
            .apply_metering_correction(&good, &actor, 10)
            .is_err());
        projection
            .run_status
            .insert(good.run_id.clone(), RunState::Terminal(RunStatus::Failed));
        let mut pending = good.metering.clone();
        if let MeteringDetails::Model {
            source, outcome, ..
        } = &mut pending
        {
            *source = MeteringSource::Estimated;
            *outcome = MeteringOutcome::Unknown;
        }
        projection.pending_usage_operations.insert(
            (good.run_id.clone(), good.operation.clone()),
            (100, 20, 0, Some(pending)),
        );
        assert_eq!(
            projection
                .operation_usage(&good.run_id, &good.operation)
                .unwrap()
                .prompt_tokens,
            100
        );
        let mut wrong = good.clone();
        if let MeteringDetails::Model { provider, .. } = &mut wrong.metering {
            *provider = Some("wrong".into());
        }
        assert!(projection
            .apply_metering_correction(&wrong, &actor, 10)
            .is_err());
        projection
            .apply_metering_correction(&good, &actor, 10)
            .unwrap();
        projection
            .apply_metering_correction(&good, &actor, 11)
            .unwrap();
        assert_eq!(projection.metering_correction(&good.id).unwrap().seq, 10);
        assert!(projection
            .apply_metering_correction(&good, &"other".parse().unwrap(), 12)
            .is_err());
        assert_eq!(
            projection
                .operation_usage(&good.run_id, &good.operation)
                .unwrap()
                .prompt_tokens,
            1
        );
        assert_eq!(projection.billable_units_for("run"), 120);
        assert!(projection
            .operation_usage(&good.run_id, "missing")
            .is_none());
    }
}