agent_engine/engine/

session.rs

//! Engine-level session management — save, load, resume, clear.
//!
//! Owns the conversation state that both TUI and headless modes need:
//! messages, token counts, cost.

use crate::pricing::calculate_cost_optional_split;
use crate::SharedMessage;
use crate::{Runtime, Session};

/// Session-scoped durability barrier state, shared by persistent frontends.
/// A failed (or dropped in-flight) publication requires explicit reload/new
/// session before saving or scheduling more inference. Never rollback-save an
/// old head: `save_durable` can fail after publishing its replacement.
#[derive(Default)]
pub struct ContextHeadPersistence {
    blocked_session: Option<String>,
}

impl ContextHeadPersistence {
    pub fn is_blocked(&self, session: &Session) -> bool {
        self.blocked_session.as_deref() == Some(session.id.as_str())
    }

    /// Candidate metadata must come from the host's current session, never
    /// from message content. The event supplies only identity and messages.
    pub async fn persist(
        &mut self,
        session: &mut Session,
        messages: &mut Vec<SharedMessage>,
        session_id: &str,
        candidate: Session,
    ) -> std::io::Result<()> {
        self.persist_with(
            session,
            messages,
            session_id,
            candidate,
            |candidate| async move { candidate.save_durable().await },
        )
        .await
    }

    async fn persist_with<F, Fut>(
        &mut self,
        session: &mut Session,
        messages: &mut Vec<SharedMessage>,
        session_id: &str,
        candidate: Session,
        save: F,
    ) -> std::io::Result<()>
    where
        F: FnOnce(Session) -> Fut,
        Fut: std::future::Future<Output = std::io::Result<()>>,
    {
        if self.is_blocked(session) {
            return Err(std::io::Error::other(
                "context head is unverified; reload the session before continuing",
            ));
        }
        // Latch before any await, including on identity rejection. Dropping
        // this future cannot authorize a later ordinary save of the old head.
        self.blocked_session = Some(session.id.clone());
        if session_id != session.id || candidate.id != session.id {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "context head checkpoint does not match the current session",
            ));
        }
        let result = save(candidate.clone()).await;
        // Success adopts the durable head. An attempted-save error adopts it
        // ONLY as conservative recovery state (not a successful commit): a
        // rename may already have happened. The latch remains set, so neither
        // post-turn saves nor automatic inference can act on this ambiguity.
        *messages = candidate.api_messages.clone();
        *session = candidate;
        if result.is_ok() {
            self.blocked_session = None;
        }
        result
    }
}

/// Conversation state tracked by the engine.
pub struct ConversationState {
    pub session: Session,
    pub context_head: ContextHeadPersistence,
    pub api_messages: Vec<SharedMessage>,
    pub total_input_tokens: u64,
    pub total_output_tokens: u64,
    pub total_cache_read_tokens: u64,
    pub total_cache_creation_tokens: u64,
    pub session_cost: f64,
    /// Prompt-cache keep-warm spend (task #444), already INCLUDED in
    /// `session_cost` and shown separately; persisted with it
    /// (`Session::keepwarm_cost`).
    pub keepwarm_cost: f64,
    /// Keep-warm ping tokens, already INCLUDED in the `total_*` counters
    /// above and shown separately so `/stats` adds up.
    pub keepwarm_tokens: crate::session::ConversationTokens,
    /// Message queued to send after current stream completes.
    pub queued_message: Option<String>,
    /// Events buffered during streaming — drained on stream completion.
    pub pending_events: Vec<String>,
}

/// The keep-warm token subtotal as persisted in the snapshot
/// (`Session::keepwarm_tokens`, snapshot-only — never journaled).
fn persisted_keepwarm_tokens(
    t: &crate::session::ConversationTokens,
) -> agent_core::core::session::KeepwarmTokens {
    agent_core::core::session::KeepwarmTokens {
        input: t.input,
        output: t.output,
        cache_read: t.cache_read,
        cache_creation: t.cache_creation,
    }
}

/// The keep-warm subtotal a resumed conversation starts with: each bucket
/// clamped under the total it is part of, so `/stats` still adds up. The
/// cache totals persist with the session (snapshot-only); a session saved
/// before they did resumes them at 0, and its keep-warm cache buckets too.
fn resumed_keepwarm_tokens(session: &Session) -> crate::session::ConversationTokens {
    let kw = session.keepwarm_tokens;
    crate::session::ConversationTokens {
        input: kw.input.min(session.total_input_tokens),
        output: kw.output.min(session.total_output_tokens),
        cache_read: kw.cache_read.min(session.total_cache_read_tokens),
        cache_creation: kw.cache_creation.min(session.total_cache_creation_tokens),
    }
}

impl ConversationState {
    /// Create a new conversation with a fresh session.
    pub fn new(session: Session) -> Self {
        Self {
            session,
            context_head: ContextHeadPersistence::default(),
            api_messages: Vec::new(),
            total_input_tokens: 0,
            total_output_tokens: 0,
            total_cache_read_tokens: 0,
            total_cache_creation_tokens: 0,
            session_cost: 0.0,
            keepwarm_cost: 0.0,
            keepwarm_tokens: Default::default(),
            queued_message: None,
            pending_events: Vec::new(),
        }
    }

    /// Create from a resumed session. A session saved before interruption
    /// markers existed (legacy `abort_context` recap) is migrated here too —
    /// `/resume` loads sessions without going through `engine::setup` —
    /// idempotently: an already-migrated session is untouched.
    pub fn from_resumed(mut session: Session) -> Self {
        crate::engine::interrupt::migrate_legacy_abort_context(
            &mut session.api_messages,
            &mut session.abort_context,
        );
        Self {
            api_messages: session.api_messages.clone(),
            total_input_tokens: session.total_input_tokens,
            total_output_tokens: session.total_output_tokens,
            total_cache_read_tokens: session.total_cache_read_tokens,
            total_cache_creation_tokens: session.total_cache_creation_tokens,
            session_cost: session.session_cost,
            keepwarm_cost: session.keepwarm_cost.min(session.session_cost).max(0.0),
            keepwarm_tokens: resumed_keepwarm_tokens(&session),
            queued_message: None,
            pending_events: Vec::new(),
            session,
            context_head: ContextHeadPersistence::default(),
        }
    }

    /// Save the current conversation state to disk. Returns whether it was
    /// written: `false` when there is nothing to save (empty history), the
    /// context head is unverified (`context_head`), or the write failed
    /// (logged).
    pub async fn save(&mut self) -> bool {
        let Some(session) = self.prepare_save() else {
            return false;
        };
        match session.save().await {
            Ok(()) => true,
            Err(e) => {
                tracing::error!("Failed to save session: {}", e);
                false
            }
        }
    }

    /// Fold the live state into `self.session` and return the snapshot to
    /// write — or `None` when there is nothing to save (empty history) or
    /// the context head is unverified. The session actor hands the snapshot
    /// to its background writer (`session::persister`) instead of awaiting
    /// the I/O. The clone is cheap: messages are `Arc`-shared.
    pub fn prepare_save(&mut self) -> Option<Session> {
        if self.context_head.is_blocked(&self.session) || self.api_messages.is_empty() {
            return None;
        }
        self.session.api_messages = self.api_messages.clone();
        self.session.total_input_tokens = self.total_input_tokens;
        self.session.total_output_tokens = self.total_output_tokens;
        self.session.total_cache_read_tokens = self.total_cache_read_tokens;
        self.session.total_cache_creation_tokens = self.total_cache_creation_tokens;
        self.session.session_cost = self.session_cost;
        self.session.keepwarm_cost = self.keepwarm_cost;
        self.session.keepwarm_tokens = persisted_keepwarm_tokens(&self.keepwarm_tokens);
        // Legacy field: an interrupted turn now lives in `api_messages`
        // (`engine::interrupt`); never write a recap back.
        self.session.abort_context = None;
        self.session.updated_at = chrono::Utc::now();
        self.session.auto_title();
        Some(self.session.clone())
    }

    /// Persist a runtime-requested head using host metadata and latest usage.
    /// The caller completes the receipt with this exact result and continues
    /// consuming the runtime's typed terminal error on failure.
    pub async fn persist_context_head(
        &mut self,
        session_id: &str,
        messages: Vec<SharedMessage>,
    ) -> std::io::Result<()> {
        let mut candidate = self.session.clone();
        candidate.api_messages = messages;
        candidate.total_input_tokens = self.total_input_tokens;
        candidate.total_output_tokens = self.total_output_tokens;
        candidate.total_cache_read_tokens = self.total_cache_read_tokens;
        candidate.total_cache_creation_tokens = self.total_cache_creation_tokens;
        candidate.session_cost = self.session_cost;
        candidate.keepwarm_cost = self.keepwarm_cost;
        candidate.keepwarm_tokens = persisted_keepwarm_tokens(&self.keepwarm_tokens);
        candidate.abort_context = None;
        candidate.updated_at = chrono::Utc::now();
        candidate.auto_title();
        self.context_head
            .persist(
                &mut self.session,
                &mut self.api_messages,
                session_id,
                candidate,
            )
            .await
    }

    /// Clear the current session and start fresh.
    pub async fn clear(&mut self, runtime: &Runtime) {
        self.save().await;
        self.context_head = ContextHeadPersistence::default();
        self.api_messages.clear();
        self.total_input_tokens = 0;
        self.total_output_tokens = 0;
        self.total_cache_read_tokens = 0;
        self.total_cache_creation_tokens = 0;
        self.session_cost = 0.0;
        self.keepwarm_cost = 0.0;
        self.keepwarm_tokens = Default::default();
        self.queued_message = None;
        self.pending_events.clear();
        self.session = Session::new(
            runtime.model(),
            runtime.thinking_level(),
            runtime.system_prompt(),
        );
        runtime.reset_context_continuation(&self.session.id, &[]);
    }

    /// Serializable mirror for clients (`SessionEventWire::Conversation`).
    /// `consecutive_auto_turns` is actor/App-side state, not conversation
    /// state, so the caller supplies it.
    pub fn snapshot(
        &self,
        consecutive_auto_turns: u32,
    ) -> crate::session::ConversationSnapshot {
        crate::session::ConversationSnapshot {
            header: crate::session::SessionHeader::from(&self.session),
            api_messages: self.api_messages.clone(),
            messages_len: self.api_messages.len(),
            tokens: crate::session::ConversationTokens {
                input: self.total_input_tokens,
                output: self.total_output_tokens,
                cache_read: self.total_cache_read_tokens,
                cache_creation: self.total_cache_creation_tokens,
            },
            cost: self.session_cost,
            keepwarm_cost: self.keepwarm_cost,
            keepwarm_tokens: self.keepwarm_tokens.clone(),
            // Wire-compat field (protocol v3); always `None` now.
            abort_context: None,
            queued_message: self.queued_message.clone(),
            pending_events_len: self.pending_events.len(),
            consecutive_auto_turns,
            messages_hash_memo: Default::default(),
        }
    }

    /// Add usage from a model turn.
    ///
    /// `cache_creation_5m` / `cache_creation_1h` are the cache-write TTL
    /// split. When either is present the cost uses the split rates
    /// (5m: 1.25×, 1h: 2.0×); when both are `None` the aggregate
    /// `cache_creation` is billed at the 5m rate (fail-cheap fallback).
    #[allow(clippy::too_many_arguments)]
    pub fn add_usage(
        &mut self,
        input_tokens: u64,
        output_tokens: u64,
        cache_read: u64,
        cache_creation: u64,
        cache_creation_5m: Option<u64>,
        cache_creation_1h: Option<u64>,
        model: &str,
    ) {
        self.total_input_tokens += input_tokens;
        self.total_output_tokens += output_tokens;
        self.total_cache_read_tokens += cache_read;
        self.total_cache_creation_tokens += cache_creation;

        self.session_cost += calculate_cost_optional_split(
            model,
            input_tokens,
            output_tokens,
            cache_read,
            cache_creation,
            cache_creation_5m,
            cache_creation_1h,
        );
    }
}

#[cfg(test)]
mod context_head_tests {
    use super::*;
    use agent_core::core::context_head::ContextHeadReceipt;
    use agent_core::core::session_journal::{save_session_durable_in_dir, SessionPersistence};
    use serde_json::json;
    use std::sync::Arc;

    fn fixture() -> (Session, Vec<SharedMessage>, Session) {
        let mut old = Session::new("synthetic-model", "medium", Some("host authority"));
        old.id = "synthetic-context-head".into();
        old.name = Some("synthetic-name".into());
        old.title = "original title".into();
        old.parent_session = Some("synthetic-parent".into());
        old.api_messages = vec![Arc::new(json!({"role":"user","content":"old head"}))];
        let messages = old.api_messages.clone();
        let mut candidate = old.clone();
        candidate.api_messages = vec![Arc::new(json!({"role":"user","content":"candidate"}))];
        candidate.total_input_tokens = 23;
        candidate.total_output_tokens = 7;
        candidate.session_cost = 0.125;
        candidate.abort_context = Some("latest synthetic abort context".into());
        (old, messages, candidate)
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn context_head_ack_follows_durable_write_and_adoption() {
        for mode in [SessionPersistence::Json, SessionPersistence::Journal] {
            let dir = tempfile::tempdir().unwrap();
            let (mut session, mut messages, candidate) = fixture();
            let id = session.id.clone();
            let expected = serde_json::to_value(&candidate).unwrap();
            let mut state = ContextHeadPersistence::default();
            let (receipt, mut acknowledged) = ContextHeadReceipt::channel();
            let (started_tx, started_rx) = tokio::sync::oneshot::channel();
            let (release_tx, release_rx) = tokio::sync::oneshot::channel();
            let path = dir.path().to_owned();
            let result = {
                let persist = state.persist_with(
                    &mut session,
                    &mut messages,
                    &id,
                    candidate,
                    |candidate| async move {
                        started_tx.send(()).unwrap();
                        release_rx.await.unwrap();
                        save_session_durable_in_dir(&path, &candidate, mode).map(|_| ())
                    },
                );
                tokio::pin!(persist);
                tokio::select! {
                    result = &mut persist => panic!("write unexpectedly completed: {result:?}"),
                    _ = started_rx => {}
                }
                assert!(matches!(
                    acknowledged.try_recv(),
                    Err(tokio::sync::oneshot::error::TryRecvError::Empty)
                ));
                release_tx.send(()).unwrap();
                persist.await
            };
            assert!(result.is_ok());
            receipt.complete(result);
            assert_eq!(acknowledged.await.unwrap(), Ok(()));
            assert_eq!(serde_json::to_value(&session).unwrap(), expected);
            assert_eq!(messages, session.api_messages);
            assert!(!state.is_blocked(&session));
            let loaded = Session::load_from_dir(dir.path(), &id).unwrap();
            assert_eq!(loaded.api_messages, messages);
            assert_eq!(loaded.system_prompt.as_deref(), Some("host authority"));
            assert_eq!(loaded.name, session.name);
            assert_eq!(loaded.total_input_tokens, 23);
            assert_eq!(loaded.total_output_tokens, 7);
            assert_eq!(loaded.abort_context, session.abort_context);
        }
    }

    #[tokio::test]
    async fn context_head_identity_mismatch_never_writes_or_adopts() {
        let (mut session, mut messages, candidate) = fixture();
        let before = serde_json::to_value(&session).unwrap();
        let mut state = ContextHeadPersistence::default();
        let (receipt, acknowledged) = ContextHeadReceipt::channel();
        let result = state
            .persist_with(
                &mut session,
                &mut messages,
                "another-session",
                candidate,
                |_| async { panic!("identity rejection must not reach persistence") },
            )
            .await;
        assert_eq!(
            result.as_ref().unwrap_err().kind(),
            std::io::ErrorKind::InvalidInput
        );
        receipt.complete(result);
        assert!(acknowledged.await.unwrap().is_err());
        assert_eq!(serde_json::to_value(&session).unwrap(), before);
        assert_eq!(messages, session.api_messages);
        assert!(state.is_blocked(&session));
        let mut other = session.clone();
        other.id = "new-session".into();
        assert!(!state.is_blocked(&other));
    }

    /// The cache totals persist with the session, so a daemon reload, an
    /// unpark or `/resume` keeps the footer's cache hit rate. The keep-warm
    /// cache buckets come back clamped under the totals they are part of.
    #[test]
    fn cache_totals_survive_save_and_resume() {
        let mut s = Session::new("synthetic-model", "low", None);
        s.api_messages = vec![Arc::new(json!({"role":"user","content":"q"}))];
        let mut conv = ConversationState::from_resumed(s);
        conv.add_usage(120, 40, 85_679, 6_250, None, None, "synthetic-model");
        conv.keepwarm_tokens.cache_read = 4_000;
        conv.keepwarm_tokens.cache_creation = 9_999; // stale label, above its total
        let saved = conv.prepare_save().expect("non-empty history saves");
        assert_eq!(saved.total_cache_read_tokens, 85_679);
        assert_eq!(saved.total_cache_creation_tokens, 6_250);

        let resumed = ConversationState::from_resumed(saved);
        assert_eq!(resumed.total_cache_read_tokens, 85_679);
        assert_eq!(resumed.total_cache_creation_tokens, 6_250);
        assert_eq!(resumed.keepwarm_tokens.cache_read, 4_000);
        assert_eq!(resumed.keepwarm_tokens.cache_creation, 6_250);
        let wire = resumed.snapshot(0);
        assert_eq!(wire.tokens.cache_read, 85_679);
        assert_eq!(wire.tokens.cache_creation, 6_250);
    }

    /// `/resume` builds its conversation with `from_resumed` directly: a
    /// legacy recap is migrated (dropped, marker appended), never carried.
    #[test]
    fn from_resumed_migrates_a_legacy_abort_context() {
        let mut legacy = Session::new("synthetic-model", "low", None);
        legacy.api_messages = vec![Arc::new(json!({"role":"user","content":"do X"}))];
        legacy.abort_context = Some("(System note — ABORT CONTEXT: …)".into());
        let conv = ConversationState::from_resumed(legacy);
        assert!(conv.session.abort_context.is_none());
        assert_eq!(conv.api_messages.len(), 2);
        assert_eq!(
            conv.api_messages[1]["content"],
            crate::engine::interrupt::InterruptReason::Unknown.marker()
        );
        assert_eq!(conv.session.api_messages, conv.api_messages);
        // Idempotent for an already-migrated session.
        let again = ConversationState::from_resumed(conv.session.clone());
        assert_eq!(again.api_messages, conv.api_messages);
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn context_head_post_publish_error_is_not_rollback_or_success() {
        let dir = tempfile::tempdir().unwrap();
        let (session, _, candidate) = fixture();
        let mut conv = ConversationState::from_resumed(session);
        let id = conv.session.id.clone();
        let expected = candidate.api_messages.clone();
        let path = dir.path().to_owned();
        let (receipt, acknowledged) = ContextHeadReceipt::channel();
        let result = conv
            .context_head
            .persist_with(
                &mut conv.session,
                &mut conv.api_messages,
                &id,
                candidate,
                |candidate| async move {
                    save_session_durable_in_dir(&path, &candidate, SessionPersistence::Journal)?;
                    Err(std::io::Error::other("synthetic error after publication"))
                },
            )
            .await;
        receipt.complete(result);
        assert!(acknowledged.await.unwrap().is_err());
        assert!(conv.context_head.is_blocked(&conv.session));
        assert_eq!(conv.api_messages, expected);
        assert_eq!(
            Session::load_from_dir(dir.path(), &id)
                .unwrap()
                .api_messages,
            expected
        );
        // Simulate stale failure repair/UI data. save() must return before it
        // even copies this old state into the conservatively adopted Session.
        conv.api_messages.clear();
        conv.api_messages
            .push(Arc::new(json!({"role":"user","content":"stale"})));
        conv.save().await;
        assert_eq!(conv.session.api_messages, expected);
        let loaded = Session::load_from_dir(dir.path(), &id).unwrap();
        let resumed = ConversationState::from_resumed(loaded);
        assert!(!resumed.context_head.is_blocked(&resumed.session));
    }

    #[tokio::test]
    async fn context_head_dropped_write_future_latches_without_mutating_old_head() {
        let (mut session, mut messages, candidate) = fixture();
        let before = serde_json::to_value(&session).unwrap();
        let id = session.id.clone();
        let mut state = ContextHeadPersistence::default();
        {
            let persist = state.persist_with(&mut session, &mut messages, &id, candidate, |_| {
                std::future::pending::<std::io::Result<()>>()
            });
            tokio::pin!(persist);
            assert!(futures::poll!(&mut persist).is_pending());
        }
        assert!(state.is_blocked(&session));
        assert_eq!(serde_json::to_value(&session).unwrap(), before);
        assert_eq!(messages, session.api_messages);
    }
}