magi-code 0.77.1

//! Session accounting is independent of the bounded transcript and activity tree.
use crate::output::NormalizedUsageSnapshot;
pub(crate) use crate::sessions::usage::{
    SessionUsageLedger, SessionUsageRecord, USAGE_ACTIVITY_PREFIX,
};

impl super::state::MissionControlState {
    pub(crate) fn session_usage_totals(&self) -> NormalizedUsageSnapshot {
        let mut usage = self.session_usage.totals();
        usage.cache_known &= self.session_usage_is_complete();
        usage
    }

    /// False for legacy/malformed history whose full request totals cannot be recovered.
    pub(crate) fn session_usage_is_complete(&self) -> bool {
        self.session_usage.is_complete()
    }

    pub(crate) fn session_cache_percent(&self) -> Option<u64> {
        let usage = self.session_usage_totals();
        if usage.cache_known && usage.effective_input > 0 {
            Some(
                (u128::from(usage.cache_read) * 100 / u128::from(usage.effective_input)).min(100)
                    as u64,
            )
        } else {
            self.last_known_session_cache_percent
        }
    }

    pub(crate) fn refresh_session_cache_percent(&mut self) {
        self.last_known_session_cache_percent = self.session_cache_percent();
    }
}

/// Called only by execution workers, before attempting bounded event delivery.
pub(crate) struct SessionUsageRecorder {
    run_id: String,
    session: Option<crate::sessions::Session>,
    cwd: std::path::PathBuf,
    pub(crate) ledger: std::sync::Arc<std::sync::Mutex<SessionUsageLedger>>,
}

impl SessionUsageRecorder {
    pub(crate) fn new(
        session: Option<crate::sessions::Session>,
        cwd: std::path::PathBuf,
        ledger: std::sync::Arc<std::sync::Mutex<SessionUsageLedger>>,
    ) -> Self {
        Self {
            run_id: uuid::Uuid::new_v4().to_string(),
            session,
            cwd,
            ledger,
        }
    }

    pub(crate) fn request_started(
        &self,
        source: &str,
        request_sequence: u64,
    ) -> anyhow::Result<Option<crate::output::ActivityEvent>> {
        let record = self
            .ledger
            .lock()
            .unwrap_or_else(|error| error.into_inner())
            .pending_request(
                format!("{USAGE_ACTIVITY_PREFIX}{}/{source}", self.run_id),
                request_sequence,
            );
        record
            .map(|record| self.record(source, record.usage, request_sequence, false))
            .transpose()
    }

    pub(crate) fn record_activity(
        &self,
        event: &crate::output::ActivityEvent,
    ) -> anyhow::Result<Option<crate::output::ActivityEvent>> {
        use crate::output::ActivityEvent;
        match event {
            ActivityEvent::UsageSnapshot {
                id,
                usage,
                request_sequence,
                final_usage,
            } => self
                .record(
                    &format!("activity/{}", id.0),
                    usage.whole_run,
                    *request_sequence,
                    *final_usage,
                )
                .map(Some),
            ActivityEvent::UsageUpdate {
                id,
                request_sequence,
                ..
            } => self.request_started(&format!("activity/{}", id.0), *request_sequence),
            _ => Ok(None),
        }
    }

    pub(crate) fn record(
        &self,
        source: &str,
        usage: NormalizedUsageSnapshot,
        request_sequence: u64,
        final_usage: bool,
    ) -> anyhow::Result<crate::output::ActivityEvent> {
        let record = SessionUsageRecord {
            id: format!("{USAGE_ACTIVITY_PREFIX}{}/{source}", self.run_id),
            usage,
            request_sequence,
            final_usage,
        };
        self.ledger
            .lock()
            .unwrap_or_else(|error| error.into_inner())
            .observe(record.clone());
        crate::sessions::record_session_event(
            self.session.as_ref(),
            &self.cwd,
            crate::sessions::SessionEventKind::SessionUsage,
            serde_json::to_value(&record)?,
        )?;
        Ok(crate::output::ActivityEvent::UsageSnapshot {
            id: crate::output::ActivityId::new(record.id),
            usage: crate::output::NormalizedUsageAggregate {
                whole_run: usage,
                ..Default::default()
            },
            request_sequence,
            final_usage,
        })
    }
}

#[cfg(test)]
mod tests;