magi-code 0.105.1

//! Best-effort diagnostic snapshot, separate from durable session history.
use super::super::Session;
use crate::{
    config::{EffectiveConfig, Settings},
    persistence::{atomic_write_with_permissions, read_regular_file_bounded},
    progressive_session_maintenance::Trigger,
};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::{path::PathBuf, time::Instant};

const MAX_STATUS_BYTES: u64 = 16 * 1024;

#[derive(Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum Stage {
    WaitingForLock,
    ReadingSummaries,
    ReadingHistory,
    PreparingProvider,
    RequestingSummary,
    ReceivingSummary,
    ValidatingSummary,
    WritingSummaries,
    WaitingForToolResult,
    Completed,
    Failed,
}

#[derive(Serialize, Deserialize)]
struct Failure {
    at: DateTime<Utc>,
    stage: Stage,
    message: String,
}

#[derive(Serialize, Deserialize)]
struct Snapshot {
    schema_version: u8,
    session_id: String,
    process_id: u32,
    started_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
    elapsed_ms: u64,
    stage: Stage,
    through_event_id: String,
    turn_ended: bool,
    provider: String,
    model: String,
    reasoning: String,
    pending_records: usize,
    batch_records: usize,
    batch_input_bytes: usize,
    saved_records: usize,
    saved_batches: usize,
    last_success_at: Option<DateTime<Utc>>,
    last_error: Option<Failure>,
}

pub(crate) struct Status {
    path: PathBuf,
    started: Instant,
    snapshot: Snapshot,
}

impl Status {
    pub(crate) fn new(
        session: &Session,
        trigger: &Trigger,
        config: &EffectiveConfig,
        settings: &Settings,
    ) -> Self {
        let path = session
            .path()
            .with_file_name(format!("{}.psm.status.json", session.id()));
        let previous = path.parent().and_then(|root| {
            let file = super::super::store::open_existing_named(root, path.file_name()?.to_str()?)
                .ok()??;
            let bytes = read_regular_file_bounded(&file, MAX_STATUS_BYTES).ok()?;
            let mut snapshot: Snapshot = serde_json::from_slice(&bytes).ok()?;
            if let Some(error) = snapshot.last_error.as_mut() {
                error.message = safe_text(&error.message, 2000);
            }
            (snapshot.schema_version == 1 && snapshot.session_id == session.id())
                .then_some(snapshot)
        });
        let now = Utc::now();
        Self {
            path,
            started: Instant::now(),
            snapshot: Snapshot {
                schema_version: 1,
                session_id: session.id().to_owned(),
                process_id: std::process::id(),
                started_at: now,
                updated_at: now,
                elapsed_ms: 0,
                stage: Stage::WaitingForLock,
                through_event_id: trigger.through_event_id.clone(),
                turn_ended: trigger.turn_ended,
                provider: safe_text(
                    settings
                        .psm
                        .provider
                        .as_deref()
                        .unwrap_or(config.provider_id()),
                    256,
                ),
                model: safe_text(
                    settings
                        .psm
                        .model
                        .as_deref()
                        .or(config.model.as_deref())
                        .unwrap_or("provider default"),
                    256,
                ),
                reasoning: settings
                    .psm
                    .reasoning
                    .unwrap_or(config.thinking_level)
                    .to_string(),
                pending_records: 0,
                batch_records: 0,
                batch_input_bytes: 0,
                saved_records: 0,
                saved_batches: 0,
                last_success_at: previous
                    .as_ref()
                    .and_then(|snapshot| snapshot.last_success_at),
                last_error: previous.and_then(|snapshot| snapshot.last_error),
            },
        }
    }

    pub(crate) fn stage(&mut self, stage: Stage) {
        self.snapshot.stage = stage;
        self.save();
    }

    pub(crate) fn selection(&mut self, selection: &crate::providers::ProviderSelection) {
        self.snapshot.provider = safe_text(&selection.provider, 256);
        self.snapshot.model = safe_text(&selection.model, 256);
    }

    pub(crate) fn pending(&mut self, records: usize) {
        self.snapshot.pending_records = records;
    }

    pub(crate) fn batch(&mut self, records: usize, bytes: usize) {
        self.snapshot.batch_records = records;
        self.snapshot.batch_input_bytes = bytes;
        self.stage(Stage::PreparingProvider);
    }

    pub(crate) fn saved(&mut self, records: usize) {
        self.snapshot.saved_records += records;
        self.snapshot.saved_batches += 1;
        self.snapshot.last_success_at = Some(Utc::now());
        self.save();
    }

    pub(crate) fn finish(&mut self, result: &anyhow::Result<bool>) {
        match result {
            Ok(true) => self.stage(Stage::WaitingForToolResult),
            Ok(false) => self.stage(Stage::Completed),
            Err(error) => {
                self.snapshot.last_error = Some(Failure {
                    at: Utc::now(),
                    stage: self.snapshot.stage,
                    message: safe_text(&format!("{error:#}"), 2000),
                });
                self.stage(Stage::Failed);
            }
        }
    }

    fn save(&mut self) {
        self.snapshot.updated_at = Utc::now();
        self.snapshot.elapsed_ms = self.started.elapsed().as_millis().min(u64::MAX as u128) as u64;
        // Diagnostics cannot fail a batch or mutate conversation revisions.
        let _ = (|| -> anyhow::Result<()> {
            let root = self
                .path
                .parent()
                .ok_or_else(|| anyhow::anyhow!("status root missing"))?;
            super::super::store::prepare_session_root(root)?;
            let name = self
                .path
                .file_name()
                .and_then(|name| name.to_str())
                .ok_or_else(|| anyhow::anyhow!("status filename missing"))?;
            // Reject unsafe existing targets with the same held-file policy as summaries.
            super::super::store::open_existing_named(root, name)?;
            let bytes = serde_json::to_vec_pretty(&self.snapshot)?;
            anyhow::ensure!(
                bytes.len() as u64 <= MAX_STATUS_BYTES,
                "PSM status exceeds limit"
            );
            atomic_write_with_permissions(
                &self.path,
                &bytes,
                Some(super::super::store::SESSION_FILE_MODE),
            )
        })();
    }
}

fn safe_text(text: &str, limit: usize) -> String {
    crate::output::redact_sensitive_text(text)
        .chars()
        .map(|character| {
            if character.is_control() {
                ' '
            } else {
                character
            }
        })
        .take(limit)
        .collect()
}