af-agent-session 0.8.0

//! Attribution attached to immutable operation usage. Monetary policy stays in the host.
use af_context::{ProviderAttemptId, ToolCallId};
use serde::{Deserialize, Serialize};

/// Origin of token or resource measurements.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MeteringSource {
    /// The executing provider supplied the measurement.
    Reported,
    /// Factory estimated usage or applied a declared resource weight.
    Estimated,
}

/// Result known at the time this immutable measurement was recorded.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MeteringOutcome {
    /// The operation returned successfully.
    Completed,
    /// The operation returned an error.
    Failed,
    /// The caller cancelled the operation.
    Cancelled,
    /// New input interrupted the operation.
    Steered,
    /// Deadline elapsed before the result was known.
    TimedOut,
    /// Admission failed before provider dispatch.
    NotDispatched,
    /// A durable preparation survived but the result did not.
    Unknown,
}

/// Stable attribution for one operation; contains no request, secret or price.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum MeteringDetails {
    /// One model invocation (including summarization).
    Model {
        /// Exact requested model.
        model: String,
        /// Registered provider; absent for an unlabelled custom adapter.
        provider: Option<String>,
        /// Stable identity recorded before dispatch.
        provider_attempt_id: ProviderAttemptId,
        /// Whether usage came from the provider.
        source: MeteringSource,
        /// Result known when recording usage.
        outcome: MeteringOutcome,
    },
    /// One tool invocation.
    Tool {
        /// Stable tool call identity.
        call_id: ToolCallId,
        /// Registered tool name.
        name: String,
        /// Declared resource weights are estimates, not prices.
        source: MeteringSource,
        /// Result known when recording usage.
        outcome: MeteringOutcome,
    },
}

impl MeteringDetails {
    /// A result must preserve the preparation's identity; unknown provider may become known.
    pub fn completes(&self, prepared: &Self) -> bool {
        match (self, prepared) {
            (
                Self::Model {
                    model,
                    provider,
                    provider_attempt_id,
                    ..
                },
                Self::Model {
                    model: expected_model,
                    provider: expected_provider,
                    provider_attempt_id: expected_attempt,
                    ..
                },
            ) => {
                model == expected_model
                    && provider_attempt_id == expected_attempt
                    && expected_provider
                        .as_ref()
                        .is_none_or(|expected| provider.as_ref() == Some(expected))
            }
            (
                Self::Tool { call_id, name, .. },
                Self::Tool {
                    call_id: expected_call,
                    name: expected_name,
                    ..
                },
            ) => call_id == expected_call && name == expected_name,
            _ => false,
        }
    }

    /// Reject malformed identifiers and unbounded attribution at the write boundary.
    pub fn validate(&self) -> Result<(), crate::EventError> {
        let valid =
            |s: &str| !s.trim().is_empty() && s.len() <= 512 && !s.chars().any(char::is_control);
        let ok = match self {
            Self::Model {
                model,
                provider,
                provider_attempt_id,
                ..
            } => {
                valid(model)
                    && provider.as_deref().is_none_or(valid)
                    && valid(provider_attempt_id.as_str())
            }
            Self::Tool { call_id, name, .. } => valid(call_id.as_str()) && valid(name),
        };
        if ok {
            Ok(())
        } else {
            Err(crate::EventError::InvalidMetering)
        }
    }
}

impl MeteringSource {
    /// Stable protocol spelling.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Reported => "reported",
            Self::Estimated => "estimated",
        }
    }
}
impl MeteringOutcome {
    /// Stable protocol spelling.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Completed => "completed",
            Self::Failed => "failed",
            Self::Cancelled => "cancelled",
            Self::Steered => "steered",
            Self::TimedOut => "timed_out",
            Self::NotDispatched => "not_dispatched",
            Self::Unknown => "unknown",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn attribution_is_bounded_and_preserves_unknown_provider_and_outcomes() {
        for outcome in [
            MeteringOutcome::Completed,
            MeteringOutcome::Failed,
            MeteringOutcome::Cancelled,
            MeteringOutcome::Steered,
            MeteringOutcome::TimedOut,
            MeteringOutcome::NotDispatched,
            MeteringOutcome::Unknown,
        ] {
            let mut details = MeteringDetails::Model {
                model: "m".into(),
                provider: None,
                provider_attempt_id: "attempt".parse().unwrap(),
                source: MeteringSource::Estimated,
                outcome,
            };
            details.validate().unwrap();
            assert_eq!(serde_json::to_value(outcome).unwrap(), outcome.as_str());
            assert_eq!(
                serde_json::from_value::<MeteringDetails>(serde_json::to_value(&details).unwrap())
                    .unwrap(),
                details
            );
            if let MeteringDetails::Model { provider, .. } = &mut details {
                *provider = Some("\n".into());
            }
            assert!(details.validate().is_err());
        }
        for name in ["".to_owned(), "a".repeat(513), "a\nb".into()] {
            assert!(MeteringDetails::Tool {
                name,
                call_id: "call".parse().unwrap(),
                source: MeteringSource::Estimated,
                outcome: MeteringOutcome::Unknown
            }
            .validate()
            .is_err());
        }
    }
}