roma_core/

session.rs

//! Chat session: message history and persistence.
//!
//! Replaces GenericAgent's `BaseSession` inheritance chain. A `ChatSession`
//! owns the conversation history, a cache-injection strategy, and a session
//! id. History compression is left to `HistoryCompressor` (Phase 3).

use std::path::Path;

use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::error::ClassifiedError;
use crate::types::{ContentBlock, Message, NormalizedResponse, Role, ToolResultMessage};

/// Where memory snapshots should be injected by [`crate::inject_memory`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CacheStrategy {
    /// Inject in the last user message; protects Anthropic prefix cache.
    #[default]
    AnthropicPrefixCache,
    /// Inject into the system prompt (OpenAI-style).
    SystemPrompt,
    /// Do not inject.
    Disabled,
}

/// Conversation state.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatSession {
    pub session_id: Uuid,
    /// Optional system prompt.
    #[serde(default)]
    pub system: Option<String>,
    /// Message history in normalized form.
    pub messages: Vec<Message>,
    /// Cache-injection strategy.
    #[serde(default)]
    pub cache_strategy: CacheStrategy,
}

impl ChatSession {
    #[must_use]
    pub fn new() -> Self {
        Self {
            session_id: Uuid::new_v4(),
            system: None,
            messages: Vec::new(),
            cache_strategy: CacheStrategy::default(),
        }
    }

    #[must_use]
    pub fn with_system(mut self, system: impl Into<String>) -> Self {
        self.system = Some(system.into());
        self
    }

    #[must_use]
    pub fn with_cache_strategy(mut self, strategy: CacheStrategy) -> Self {
        self.cache_strategy = strategy;
        self
    }

    /// Append a user text message.
    pub fn push_user(&mut self, content: impl Into<String>) {
        self.messages.push(Message::user_text(content));
    }

    /// Append the assistant's response content blocks verbatim.
    pub fn push_assistant(&mut self, response: NormalizedResponse) {
        self.messages.push(Message {
            role: Role::Assistant,
            content: response.content_blocks,
        });
    }

    /// Append tool results as a single user message containing
    /// [`ContentBlock::ToolResult`] blocks. Providers translate the exact
    /// shape at the wire layer.
    pub fn push_tool_results(&mut self, results: Vec<ToolResultMessage>) {
        if results.is_empty() {
            return;
        }
        let content = results.into_iter().map(ContentBlock::from).collect();
        self.messages.push(Message {
            role: Role::User,
            content,
        });
    }

    /// Append text to the tail of the most recent user message as a new
    /// [`ContentBlock::Text`] block. Used for mid-turn injections (stream
    /// rules) that must NOT create a bare user message after tool results
    /// (provider alternation rules — spec D7). Falls back to pushing a new
    /// user message when no user message exists; that only happens with an
    /// empty, system-only, or assistant-only history, where a user message
    /// is alternation-safe.
    ///
    /// Caller contract: intended for mid-turn injection while the tail
    /// message is a user message (before any assistant commit). If the
    /// tail is an assistant message, the text lands on an older user
    /// message mid-history.
    pub fn append_to_last_user(&mut self, text: impl Into<String>) {
        let text = text.into();
        if text.is_empty() {
            return;
        }
        match self.messages.iter_mut().rfind(|m| m.role == Role::User) {
            Some(msg) => msg.content.push(ContentBlock::Text { text }),
            None => self.messages.push(Message::user_text(text)),
        }
    }

    /// Total message count (including system messages if pushed manually).
    #[must_use]
    pub fn len(&self) -> usize {
        self.messages.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    /// Compress history in place to fit within `target_tokens` using the
    /// provided compressor and tokenizer. Returns the token count after
    /// compression (may still exceed `target_tokens` if no further moves
    /// were safe).
    pub fn compress(
        &mut self,
        target_tokens: u32,
        compressor: &dyn crate::HistoryCompressor,
        tokenizer: &dyn crate::Tokenizer,
    ) -> u32 {
        compressor.compress(&mut self.messages, target_tokens, tokenizer)
    }

    /// Persist the session to disk as JSON (supports `/resume`).
    pub fn save(&self, path: &Path) -> Result<(), ClassifiedError> {
        let json = serde_json::to_string_pretty(self)
            .map_err(|e| ClassifiedError::Parse(format!("save session: {e}")))?;
        let tmp = path.with_extension(format!("tmp.{}", uuid::Uuid::new_v4()));
        std::fs::write(&tmp, &json)?;
        std::fs::rename(&tmp, path)?;
        Ok(())
    }

    /// Load a session from disk.
    pub fn load(path: &Path) -> Result<Self, ClassifiedError> {
        let text = std::fs::read_to_string(path)?;
        serde_json::from_str(&text)
            .map_err(|e| ClassifiedError::Parse(format!("load session: {e}")))
    }
}

impl Default for ChatSession {
    fn default() -> Self {
        Self::new()
    }
}

/// Inject a memory snapshot into the session according to the configured
/// [`CacheStrategy`].
///
/// - **AnthropicPrefixCache**: Prepend to the last user message (or insert a
///   new one if none exists). This preserves the Anthropic prefix cache — the
///   system prompt and all prior turns remain intact — while avoiding
///   consecutive user messages that violate Anthropic's turn-alternation rule.
/// - **SystemPrompt**: Append to the system prompt.
/// - **Disabled**: No-op.
pub fn inject_memory(session: &mut ChatSession, snapshot: &str) {
    match session.cache_strategy {
        CacheStrategy::AnthropicPrefixCache => {
            let last_user = session.messages.iter_mut().rfind(|m| m.role == Role::User);
            match last_user {
                Some(msg) => {
                    // If the last user message contains only ToolResult blocks,
                    // prepend a new Text block with memory. This avoids injecting
                    // text into a message that is semantically tool-result-only.
                    let has_text = msg
                        .content
                        .iter()
                        .any(|b| matches!(b, ContentBlock::Text { .. }));
                    if !has_text {
                        msg.content.insert(
                            0,
                            ContentBlock::Text {
                                text: format!("[memory]\n{snapshot}"),
                            },
                        );
                    } else {
                        // Prepend memory into the first Text block to avoid
                        // consecutive user messages (Anthropic requires
                        // strict user/assistant alternation). Preserve any
                        // ToolResult blocks so that tool_use/tool_result pairing
                        // stays intact.
                        let mut new_blocks: Vec<ContentBlock> = Vec::new();
                        let mut text_injected = false;
                        for block in msg.content.drain(..) {
                            match block {
                                ContentBlock::Text { text } => {
                                    if text_injected {
                                        new_blocks.push(ContentBlock::Text { text });
                                    } else {
                                        text_injected = true;
                                        new_blocks.push(ContentBlock::Text {
                                            text: format!("[memory]\n{snapshot}\n\n{text}"),
                                        });
                                    }
                                }
                                other => new_blocks.push(other),
                            }
                        }
                        msg.content = new_blocks;
                    }
                }
                None => {
                    session
                        .messages
                        .push(Message::user_text(format!("[memory]\n{snapshot}")));
                }
            }
        }
        CacheStrategy::SystemPrompt => match &mut session.system {
            Some(sys) => {
                sys.push_str("\n\n");
                sys.push_str(snapshot);
            }
            None => {
                session.system = Some(snapshot.to_string());
            }
        },
        CacheStrategy::Disabled => {}
    }
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::types::{StopReason, Usage};

    #[test]
    fn push_user_appends_text_message() {
        let mut s = ChatSession::new();
        s.push_user("hello");
        assert_eq!(s.len(), 1);
        assert_eq!(s.messages[0].role, Role::User);
    }

    #[test]
    fn push_assistant_preserves_content_blocks() {
        let mut s = ChatSession::new();
        s.push_assistant(NormalizedResponse {
            content_blocks: vec![
                ContentBlock::Text { text: "ok".into() },
                ContentBlock::ToolUse {
                    id: "t1".into(),
                    name: "file_read".into(),
                    input: serde_json::json!({}),
                },
            ],
            stop_reason: StopReason::ToolUse,
            usage: Some(Usage::default()),
        });
        assert_eq!(s.messages[0].role, Role::Assistant);
        assert_eq!(s.messages[0].content.len(), 2);
    }

    #[test]
    fn push_tool_results_emits_single_user_message_with_blocks() {
        let mut s = ChatSession::new();
        s.push_tool_results(vec![
            ToolResultMessage {
                tool_use_id: "t1".into(),
                content: "a".into(),
                is_error: false,
            },
            ToolResultMessage {
                tool_use_id: "t2".into(),
                content: "b".into(),
                is_error: true,
            },
        ]);
        assert_eq!(s.len(), 1);
        assert_eq!(s.messages[0].role, Role::User);
        assert_eq!(s.messages[0].content.len(), 2);
    }

    #[test]
    fn empty_tool_results_is_noop() {
        let mut s = ChatSession::new();
        s.push_tool_results(vec![]);
        assert!(s.is_empty());
    }

    #[test]
    fn save_and_load_roundtrips() {
        let mut s = ChatSession::new()
            .with_system("sys")
            .with_cache_strategy(CacheStrategy::SystemPrompt);
        s.push_user("hi");
        let tmp = tempdir_path("roma_session_test.json");
        s.save(&tmp).unwrap();
        let loaded = ChatSession::load(&tmp).unwrap();
        assert_eq!(loaded.session_id, s.session_id);
        assert_eq!(loaded.system.as_deref(), Some("sys"));
        assert_eq!(loaded.cache_strategy, CacheStrategy::SystemPrompt);
        assert_eq!(loaded.messages.len(), 1);
        let _ = std::fs::remove_file(&tmp);
    }

    #[test]
    fn inject_anthropic_prefix_prepends_to_last_user() {
        let mut s = ChatSession::new().with_cache_strategy(CacheStrategy::AnthropicPrefixCache);
        s.push_user("first");
        s.push_assistant(NormalizedResponse {
            content_blocks: vec![ContentBlock::Text { text: "hi".into() }],
            stop_reason: StopReason::EndTurn,
            usage: None,
        });
        s.push_user("latest");
        inject_memory(&mut s, "[mem] insight");
        // Memory should be prepended into the last user message (not a new message).
        assert_eq!(s.messages.len(), 3);
        assert_eq!(s.messages[2].role, Role::User);
        let text = s.messages[2].text_content();
        assert!(text.starts_with("[memory]\n[mem] insight"));
        assert!(text.ends_with("latest"));
    }

    #[test]
    fn inject_system_prompt_appends_to_system() {
        let mut s = ChatSession::new()
            .with_system("base")
            .with_cache_strategy(CacheStrategy::SystemPrompt);
        inject_memory(&mut s, "[mem] insight");
        assert_eq!(s.system.as_deref(), Some("base\n\n[mem] insight"));
    }

    #[test]
    fn inject_disabled_is_noop() {
        let mut s = ChatSession::new().with_cache_strategy(CacheStrategy::Disabled);
        s.push_user("hi");
        inject_memory(&mut s, "[mem] insight");
        assert_eq!(s.len(), 1);
    }

    #[test]
    fn multi_turn_save_load_roundtrip_preserves_tool_results() {
        let mut s = ChatSession::new()
            .with_system("sys")
            .with_cache_strategy(CacheStrategy::SystemPrompt);
        s.push_user("read foo.rs");
        s.push_assistant(NormalizedResponse {
            content_blocks: vec![
                ContentBlock::Text {
                    text: "Reading file.".into(),
                },
                ContentBlock::ToolUse {
                    id: "t1".into(),
                    name: "file_read".into(),
                    input: serde_json::json!({"path": "foo.rs"}),
                },
            ],
            stop_reason: StopReason::ToolUse,
            usage: Some(Usage {
                input_tokens: 10,
                output_tokens: 5,
                cache_read_tokens: 0,
                cache_write_tokens: 0,
            }),
        });
        s.push_tool_results(vec![ToolResultMessage {
            tool_use_id: "t1".into(),
            content: "fn main() {}".into(),
            is_error: false,
        }]);
        s.push_assistant(NormalizedResponse {
            content_blocks: vec![ContentBlock::Text {
                text: "The file contains a main function.".into(),
            }],
            stop_reason: StopReason::EndTurn,
            usage: None,
        });

        let tmp = tempdir_path("roma_multi_turn_test.json");
        s.save(&tmp).unwrap();
        let loaded = ChatSession::load(&tmp).unwrap();
        let _ = std::fs::remove_file(&tmp);

        assert_eq!(loaded.session_id, s.session_id);
        assert_eq!(loaded.system.as_deref(), Some("sys"));
        assert_eq!(loaded.cache_strategy, CacheStrategy::SystemPrompt);
        assert_eq!(loaded.messages.len(), 4);
        assert_eq!(loaded.messages[0].role, Role::User);
        assert_eq!(loaded.messages[1].role, Role::Assistant);
        assert_eq!(loaded.messages[2].role, Role::User);
        assert!(matches!(
            &loaded.messages[2].content[0],
            ContentBlock::ToolResult { .. }
        ));
        assert_eq!(loaded.messages[3].role, Role::Assistant);
    }

    #[test]
    fn save_resume_continues_conversation() {
        let mut s = ChatSession::new();
        s.push_user("first question");

        let tmp = tempdir_path("roma_resume_test.json");
        s.save(&tmp).unwrap();
        let mut loaded = ChatSession::load(&tmp).unwrap();

        loaded.push_assistant(NormalizedResponse {
            content_blocks: vec![ContentBlock::Text {
                text: "first answer".into(),
            }],
            stop_reason: StopReason::EndTurn,
            usage: None,
        });
        loaded.push_user("second question");
        loaded.save(&tmp).unwrap();

        let resumed = ChatSession::load(&tmp).unwrap();
        let _ = std::fs::remove_file(&tmp);

        assert_eq!(resumed.session_id, s.session_id);
        assert_eq!(resumed.messages.len(), 3);
        assert_eq!(resumed.messages[0].text_content(), "first question");
        assert_eq!(resumed.messages[2].text_content(), "second question");
    }

    #[test]
    fn inject_memory_after_load_preserves_history() {
        let mut s = ChatSession::new()
            .with_system("base")
            .with_cache_strategy(CacheStrategy::SystemPrompt);
        s.push_user("hello");

        let tmp = tempdir_path("roma_inject_test.json");
        s.save(&tmp).unwrap();
        let mut loaded = ChatSession::load(&tmp).unwrap();
        let _ = std::fs::remove_file(&tmp);

        inject_memory(&mut loaded, "[mem] insight");
        assert!(loaded.system.as_deref().unwrap().contains("[mem] insight"));
        assert_eq!(loaded.messages.len(), 1, "no new messages from inject");
    }

    #[test]
    fn append_to_last_user_appends_text_block_to_newest_user_message() {
        let mut s = ChatSession::new();
        s.push_user("first");
        s.push_assistant(NormalizedResponse {
            content_blocks: vec![ContentBlock::Text { text: "ok".into() }],
            stop_reason: StopReason::EndTurn,
            usage: None,
        });
        s.push_user("second");
        s.append_to_last_user("\n\n[stream rule `r` fired]\n\nbody");
        assert_eq!(s.len(), 3, "no new message may be created");
        let last = &s.messages[2];
        assert_eq!(last.content.len(), 2);
        assert!(
            last.text_content()
                .ends_with("[stream rule `r` fired]\n\nbody")
        );
        assert_eq!(s.messages[0].text_content(), "first");
    }

    #[test]
    fn append_to_last_user_works_on_tool_result_only_message() {
        let mut s = ChatSession::new();
        s.push_tool_results(vec![ToolResultMessage {
            tool_use_id: "t1".into(),
            content: "out".into(),
            is_error: false,
        }]);
        s.append_to_last_user("note");
        assert_eq!(s.len(), 1);
        assert_eq!(s.messages[0].content.len(), 2);
        assert!(matches!(
            s.messages[0].content[1],
            ContentBlock::Text { .. }
        ));
    }

    #[test]
    fn append_to_last_user_falls_back_to_new_message_on_empty_history() {
        let mut s = ChatSession::new();
        s.append_to_last_user("note");
        assert_eq!(s.len(), 1);
        assert_eq!(s.messages[0].role, Role::User);
        assert_eq!(s.messages[0].text_content(), "note");
    }

    fn tempdir_path(name: &str) -> std::path::PathBuf {
        let mut p = std::env::temp_dir();
        p.push(name);
        p
    }
}