magi-code 0.77.1

Repository-aware CLI coding agent for terminal work
Documentation
use super::{Session, SessionEvent, SessionEventKind, SessionManager};

/// Uses the same durable user-input count as file checkpoints, including compaction.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct RewindPrompt {
    pub(crate) turn: u64,
    pub(crate) text: String,
}

fn prompt_indices(events: &[SessionEvent]) -> Vec<(usize, u64)> {
    let mut turn = 0u64;
    let mut prompts = Vec::new();
    for (index, event) in events.iter().enumerate() {
        match event.kind() {
            Some(SessionEventKind::Compaction) => {
                if let Some(count) = event
                    .payload
                    .pointer("/aggregate/user_input_count")
                    .and_then(serde_json::Value::as_u64)
                {
                    turn = count;
                }
            }
            Some(SessionEventKind::UserInput) => {
                turn = turn.saturating_add(1);
                if event
                    .payload
                    .get("auto_recovery")
                    .and_then(serde_json::Value::as_bool)
                    != Some(true)
                {
                    prompts.push((index, turn));
                }
            }
            _ => {}
        }
    }
    prompts
}

fn prompt_text(event: &SessionEvent) -> anyhow::Result<&str> {
    event
        .payload
        .get("text")
        .and_then(serde_json::Value::as_str)
        .ok_or_else(|| anyhow::anyhow!("user prompt is missing text"))
}

impl Session {
    fn rewind_events(&self) -> anyhow::Result<Vec<SessionEvent>> {
        let append_lock = super::write::session_append_lock(self.path())?;
        let _guard = append_lock
            .lock()
            .map_err(|_| anyhow::anyhow!("session append lock was poisoned"))?;
        let _file_guard = crate::persistence::CrossProcessFileLock::acquire(self.path())?;
        let read =
            self.read_events_tolerant_bounded(super::MAX_METADATA_VISIT_LINES, 32 * 1024 * 1024)?;
        anyhow::ensure!(
            read.diagnostics.is_empty(),
            "cannot rewind unreadable session history"
        );
        anyhow::ensure!(
            read.events
                .iter()
                .all(|event| event.session_id == self.id()),
            "cannot rewind history containing foreign session records"
        );
        let (_, diagnostics) =
            super::latest_valid_compaction_checkpoint_for_replay(self.id(), &read.events);
        anyhow::ensure!(
            diagnostics.is_empty(),
            "cannot rewind malformed compaction history"
        );
        Ok(read.events)
    }

    pub(crate) fn rewind_prompts(&self) -> anyhow::Result<Vec<RewindPrompt>> {
        let events = self.rewind_events()?;
        prompt_indices(&events)
            .into_iter()
            .map(|(index, turn)| {
                Ok(RewindPrompt {
                    turn,
                    text: prompt_text(&events[index])?.to_owned(),
                })
            })
            .collect()
    }

    /// Fork strictly before the selected prompt; never rewrite the source session.
    /// Event-count compaction cutoffs remain valid because the prefix keeps its order.
    pub(crate) fn fork_before_prompt(
        &self,
        manager: &SessionManager,
        target: u64,
    ) -> anyhow::Result<Session> {
        let mut events = self.rewind_events()?;
        let index = prompt_indices(&events)
            .into_iter()
            .find_map(|(index, turn)| (turn == target).then_some(index))
            .ok_or_else(|| anyhow::anyhow!("prompt target is not in retained session history"))?;
        let prompt = prompt_text(&events[index])?.to_owned();
        let cwd = events[index].cwd.clone();
        events.truncate(index);
        let fork = manager.create()?.admit_standalone_writer()?;
        for event in &mut events {
            event.session_id = fork.id().to_owned();
        }
        events.push(SessionEvent::new_kind(
            SessionEventKind::Diagnostic,
            fork.id().to_owned(),
            cwd,
            serde_json::json!({"rewind_source_session": self.id(), "before_user_turn": target, "rewind_prompt": prompt}),
        ));
        fork.append_owned_batch(events).map_err(|error| {
            anyhow::anyhow!("could not persist rewind fork {}: {}", fork.id(), error)
        })?;
        Ok(fork)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{providers::ProviderConversationItem, sessions::record_session_event};
    use serde_json::json;

    #[test]
    fn rewind_first_prompt_without_file_changes_starts_empty_and_keeps_original() {
        let temp = tempfile::TempDir::new().unwrap();
        let manager = SessionManager::new(temp.path().join("sessions"));
        let source = manager.create().unwrap();
        record_session_event(
            Some(&source),
            temp.path(),
            SessionEventKind::UserInput,
            json!({"text": "first prompt"}),
        )
        .unwrap();
        let original = std::fs::read(source.path()).unwrap();
        let resumed = manager.open_existing(source.id()).unwrap();
        assert_eq!(resumed.rewind_prompts().unwrap()[0].text, "first prompt");
        let fork = resumed.fork_before_prompt(&manager, 1).unwrap();
        assert!(
            crate::context::build_conversation_replay(Some(&fork))
                .unwrap()
                .items
                .is_empty()
        );
        assert_eq!(super::super::session_user_input_count(&fork).unwrap(), 0);
        assert_eq!(std::fs::read(source.path()).unwrap(), original);
        assert!(resumed.fork_before_prompt(&manager, 0).is_err());
        assert!(resumed.fork_before_prompt(&manager, 2).is_err());
    }

    #[test]
    fn rewind_after_compaction_keeps_summary_and_durable_prompt_numbers() {
        let temp = tempfile::TempDir::new().unwrap();
        let manager = SessionManager::new(temp.path().join("sessions"));
        let source = manager.create().unwrap();
        for (kind, text) in [
            (SessionEventKind::UserInput, "old prompt"),
            (SessionEventKind::AssistantOutput, "old answer"),
        ] {
            record_session_event(Some(&source), temp.path(), kind, json!({"text": text})).unwrap();
        }
        crate::sessions::record_session_compaction(
            &source,
            temp.path(),
            "earlier summary",
            "test",
            "model",
            2,
        )
        .unwrap();
        for text in ["retained prompt", "selected prompt"] {
            record_session_event(
                Some(&source),
                temp.path(),
                SessionEventKind::UserInput,
                json!({"text": text}),
            )
            .unwrap();
            record_session_event(
                Some(&source),
                temp.path(),
                SessionEventKind::AssistantOutput,
                json!({"text": "answer"}),
            )
            .unwrap();
        }
        let resumed = manager.open_existing(source.id()).unwrap();
        assert_eq!(
            resumed
                .rewind_prompts()
                .unwrap()
                .iter()
                .map(|prompt| prompt.turn)
                .collect::<Vec<_>>(),
            vec![2, 3]
        );
        assert!(resumed.fork_before_prompt(&manager, 1).is_err());
        let fork = resumed.fork_before_prompt(&manager, 3).unwrap();
        assert_eq!(crate::sessions::session_user_input_count(&fork).unwrap(), 2);
        let replay = crate::context::build_conversation_replay(Some(&fork)).unwrap();
        let text = replay
            .items
            .iter()
            .filter_map(|item| match item {
                ProviderConversationItem::Message(message) => Some(message.content.as_str()),
                _ => None,
            })
            .collect::<Vec<_>>()
            .join("\n");
        assert!(text.contains("earlier summary"));
        assert!(text.contains("retained prompt"));
        assert!(!text.contains("selected prompt"));
        assert!(!text.contains("old answer"));
    }
}