af-agent-session 0.8.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,
}

/// An original operation used in reservation settlement, never a corrected price.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RunOperationUsage {
    /// Stable operation identity within the Run, suitable for host deduplication.
    pub operation_id: String,
    /// Original immutable measurement, including provider/model or tool attribution.
    pub usage: OperationUsage,
    /// Preparation has no durable result; the usage is a conservative estimate.
    pub pending: bool,
}

/// 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 {
    /// Original settlement evidence in stable operation order, including pending estimates.
    /// Corrections are deliberately separate (`operation_usage`); they must never
    /// silently reprice an already settled reservation. Legacy facts lack attribution.
    /// Read through an authorized AgentStore projection; this value grants no access.
    pub fn run_usage_operations(&self, run_id: &RunId) -> Vec<RunOperationUsage> {
        let mut result: std::collections::BTreeMap<String, RunOperationUsage> = self
            .pending_usage_operations
            .iter()
            .filter(|((id, _), _)| id == run_id)
            .map(|((_, operation), (prompt, completion, cost, metering))| {
                (
                    operation.clone(),
                    RunOperationUsage {
                        operation_id: operation.clone(),
                        usage: OperationUsage {
                            prompt_tokens: *prompt,
                            completion_tokens: *completion,
                            cost_units: *cost,
                            metering: metering.clone(),
                            revision: 0,
                        },
                        pending: true,
                    },
                )
            })
            .collect();
        for ((id, operation), (prompt, completion, cost, metering)) in &self.usage_operations {
            if id == run_id {
                result.insert(
                    operation.clone(),
                    RunOperationUsage {
                        operation_id: operation.clone(),
                        usage: OperationUsage {
                            prompt_tokens: *prompt,
                            completion_tokens: *completion,
                            cost_units: *cost,
                            metering: metering.clone(),
                            revision: 0,
                        },
                        pending: false,
                    },
                );
            }
        }
        result.into_values().collect()
    }

    /// 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};

    #[test]
    fn original_evidence_keeps_mixed_models_pending_tools_and_excludes_corrections() {
        let mut projection = SessionProjection::default();
        let run: RunId = "run".parse().unwrap();
        for (operation, model) in [("model-a", "model-one"), ("model-b", "model-two")] {
            projection.usage_operations.insert(
                (run.clone(), operation.into()),
                (
                    3,
                    2,
                    0,
                    Some(MeteringDetails::Model {
                        model: model.into(),
                        provider: Some("provider".into()),
                        provider_attempt_id: operation.parse().unwrap(),
                        source: MeteringSource::Reported,
                        outcome: MeteringOutcome::Completed,
                    }),
                ),
            );
        }
        projection.pending_usage_operations.insert(
            (run.clone(), "tool".into()),
            (
                0,
                0,
                7,
                Some(MeteringDetails::Tool {
                    call_id: "call".parse().unwrap(),
                    name: "lookup".into(),
                    source: MeteringSource::Estimated,
                    outcome: MeteringOutcome::Unknown,
                }),
            ),
        );
        projection.usage_operations.insert(
            ("another-run".parse().unwrap(), "model-a".into()),
            (99, 0, 0, None),
        );
        let original = projection.run_usage_operations(&run);
        assert_eq!(original.len(), 3);
        assert!(!original[0].pending);
        assert!(original[2].pending);
        assert_ne!(original[0].usage.metering, original[1].usage.metering);
        assert_eq!(
            original
                .iter()
                .map(|fact| fact.usage.prompt_tokens
                    + fact.usage.completion_tokens
                    + fact.usage.cost_units)
                .sum::<u64>(),
            projection.billable_units_for("run")
        );
        projection.corrected_usage.insert(
            (run.clone(), "model-a".into()),
            OperationUsage {
                prompt_tokens: 1,
                completion_tokens: 1,
                cost_units: 0,
                metering: None,
                revision: 1,
            },
        );
        assert_eq!(projection.run_usage_operations(&run), original);
        assert_eq!(projection.clone().run_usage_operations(&run), original);
        assert!(projection
            .run_usage_operations(&"missing".parse().unwrap())
            .is_empty());
    }

    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());
    }
}