zagens_core/

session.rs

//! Session state management for the core engine.
//!
//! Tracks conversation history, token usage, and session metadata.

use crate::approval::ApprovalMode;
use crate::chat::{Message, SystemPrompt};
use crate::cycle::CycleBriefing;
use crate::engine::context::extract_compaction_summary_prompt;
use crate::models::Usage;
use crate::project_context::{ProjectContext, load_project_context_with_parents};
use crate::working_set::WorkingSet;
use chrono::{DateTime, Utc};
use std::path::PathBuf;

/// Session state for the engine.
#[derive(Debug, Clone)]
pub struct Session {
    /// Model being used
    pub model: String,

    /// Reasoning-effort tier for DeepSeek thinking mode:
    /// `"off" | "low" | "medium" | "high" | "max"`. `None` lets the provider
    /// apply its own defaults.
    pub reasoning_effort: Option<String>,
    /// Whether the user selected automatic reasoning effort.
    pub reasoning_effort_auto: bool,

    /// Whether the user selected automatic model routing.
    pub auto_model: bool,

    /// Workspace directory
    pub workspace: PathBuf,

    /// System prompt (optional)
    pub system_prompt: Option<SystemPrompt>,
    /// Hash of the last assembled stable system prompt. Used to avoid
    /// replacing `system_prompt` when unchanged.
    pub last_system_prompt_hash: Option<u64>,
    /// Persisted summary blocks generated by context compaction.
    pub compaction_summary_prompt: Option<SystemPrompt>,

    /// Conversation history (API format)
    pub messages: Vec<Message>,

    /// Total tokens used in this session
    pub total_usage: SessionUsage,

    /// Whether shell execution is allowed
    pub allow_shell: bool,

    /// Whether to trust paths outside workspace
    pub trust_mode: bool,

    /// Whether the current session should auto-approve tool safety checks.
    pub auto_approve: bool,

    /// Live UI approval policy used to steer the system prompt.
    pub approval_mode: ApprovalMode,

    /// Notes file path
    pub notes_path: PathBuf,

    /// MCP config path
    pub mcp_config_path: PathBuf,

    /// Session ID (for tracking)
    pub id: String,

    /// Project context loaded from AGENTS.md, etc.
    pub project_context: Option<ProjectContext>,

    /// Repo-aware working set for context management.
    pub working_set: WorkingSet,

    /// Number of cycle boundaries crossed in this session (issue #124). The
    /// active cycle index is `cycle_count + 1` (cycles are 1-based for users).
    pub cycle_count: u32,

    /// UTC start time of the *current* cycle. Updated when the engine resets
    /// the conversation buffer. Used by archive headers and the `/cycles`
    /// command's display.
    pub current_cycle_started: DateTime<Utc>,

    /// Briefings produced at past cycle boundaries, in chronological order.
    /// Bounded growth: one entry per cycle, briefing capped at ~3,000 tokens.
    pub cycle_briefings: Vec<CycleBriefing>,

    /// Provider-reported `usage.input_tokens` from the **most recent API round**
    /// (overwrite per round — not summed across tool-call rounds). Authoritative
    /// for “context size at last inference” per DeepSeek API docs.
    pub last_api_input_tokens: Option<u32>,

    /// Per-session sampling overrides (optional). `None` → provider / model defaults.
    pub temperature: Option<f32>,
    pub top_p: Option<f32>,
    /// Max output tokens for API requests; when unset, model-specific default applies.
    pub max_output_tokens: Option<u32>,
}

impl Session {
    /// Record per-round API usage. Turn totals still sum via `Turn::add_usage`.
    pub fn record_api_round_usage(&mut self, usage: &Usage) {
        if usage.input_tokens > 0 {
            self.last_api_input_tokens = Some(usage.input_tokens);
        }
    }
}

/// Cumulative usage statistics for a session.
#[derive(Debug, Clone, Default)]
#[allow(clippy::struct_field_names)]
pub struct SessionUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    #[allow(dead_code)]
    pub cache_creation_input_tokens: u64,
    #[allow(dead_code)]
    pub cache_read_input_tokens: u64,
}

impl SessionUsage {
    /// Add usage from a turn
    pub fn add(&mut self, usage: &Usage) {
        self.input_tokens += u64::from(usage.input_tokens);
        self.output_tokens += u64::from(usage.output_tokens);
        if let Some(tokens) = usage.prompt_cache_miss_tokens {
            self.cache_creation_input_tokens += u64::from(tokens);
        }
        if let Some(tokens) = usage.prompt_cache_hit_tokens {
            self.cache_read_input_tokens += u64::from(tokens);
        }
    }
}

impl Session {
    /// Create a new session
    pub fn new(
        model: String,
        workspace: PathBuf,
        allow_shell: bool,
        trust_mode: bool,
        notes_path: PathBuf,
        mcp_config_path: PathBuf,
    ) -> Self {
        // Load project context from AGENTS.md, CLAUDE.md, etc.
        let project_context = load_project_context_with_parents(&workspace);
        let has_context = project_context.has_instructions();

        Self {
            model,
            reasoning_effort: None,
            reasoning_effort_auto: false,
            auto_model: false,
            workspace,
            system_prompt: None,
            compaction_summary_prompt: None,
            messages: Vec::new(),
            total_usage: SessionUsage::default(),
            allow_shell,
            trust_mode,
            auto_approve: false,
            approval_mode: ApprovalMode::Suggest,
            notes_path,
            mcp_config_path,
            id: uuid::Uuid::new_v4().to_string(),
            project_context: if has_context {
                Some(project_context)
            } else {
                None
            },
            last_system_prompt_hash: None,
            working_set: WorkingSet::default(),
            cycle_count: 0,
            current_cycle_started: Utc::now(),
            cycle_briefings: Vec::new(),
            last_api_input_tokens: None,
            temperature: None,
            top_p: None,
            max_output_tokens: None,
        }
    }

    /// Add a message to the conversation
    pub fn add_message(&mut self, message: Message) {
        self.messages.push(message);
    }

    /// Rebuild the working set from current messages (best effort).
    pub fn rebuild_working_set(&mut self) {
        self.working_set
            .rebuild_from_messages(&self.messages, &self.workspace);
    }
}

/// Whether the user selected automatic model routing (`"auto"` label).
#[must_use]
pub fn is_auto_model_label(model: &str) -> bool {
    model.trim().eq_ignore_ascii_case("auto")
}

/// Apply runtime model selection to session state and a mirrored config field.
pub fn apply_model_selection(session: &mut Session, config_model: &mut String, model: String) {
    session.auto_model = is_auto_model_label(&model);
    session.model = model;
    config_model.clone_from(&session.model);
}

/// Apply runtime thread sync payload (messages, prompts, model, workspace).
pub fn apply_sync_session_payload(
    session: &mut Session,
    config_workspace: &mut PathBuf,
    config_model: &mut String,
    messages: Vec<Message>,
    system_prompt: Option<SystemPrompt>,
    model: String,
    workspace: PathBuf,
) {
    session.messages = messages;
    session.compaction_summary_prompt = extract_compaction_summary_prompt(system_prompt.clone());
    session.system_prompt = system_prompt;
    apply_model_selection(session, config_model, model);
    session.workspace = workspace.clone();
    *config_workspace = workspace.clone();
    let ctx = load_project_context_with_parents(&workspace);
    session.project_context = if ctx.has_instructions() {
        Some(ctx)
    } else {
        None
    };
    session.rebuild_working_set();
}

/// Index of the last `user` role message in API-format history.
#[must_use]
pub fn index_of_last_user_message(messages: &[Message]) -> Option<usize> {
    messages
        .iter()
        .enumerate()
        .rev()
        .find_map(|(idx, msg)| (msg.role == "user").then_some(idx))
}

/// Remove the last user message and everything after it (#383 `/edit`).
#[must_use]
pub fn truncate_before_last_user_message(messages: &mut Vec<Message>) -> bool {
    index_of_last_user_message(messages).is_some_and(|idx| {
        messages.truncate(idx);
        true
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::chat::ContentBlock;
    use std::path::PathBuf;

    #[test]
    fn is_auto_model_label_matches_auto_case_insensitive() {
        assert!(is_auto_model_label("auto"));
        assert!(is_auto_model_label(" Auto "));
        assert!(!is_auto_model_label("deepseek-v4-pro"));
    }

    #[test]
    fn apply_model_selection_updates_session_and_config() {
        let mut session = Session::new(
            "old".into(),
            PathBuf::from("/tmp"),
            false,
            false,
            PathBuf::from("/tmp/notes"),
            PathBuf::from("/tmp/mcp"),
        );
        let mut config_model = "old".to_string();
        apply_model_selection(&mut session, &mut config_model, "auto".into());
        assert!(session.auto_model);
        assert_eq!(session.model, "auto");
        assert_eq!(config_model, "auto");

        apply_model_selection(&mut session, &mut config_model, "deepseek-v4-pro".into());
        assert!(!session.auto_model);
        assert_eq!(session.model, "deepseek-v4-pro");
        assert_eq!(config_model, "deepseek-v4-pro");
    }

    #[test]
    fn apply_sync_session_payload_updates_messages_workspace_and_model() {
        let tmpdir = tempfile::TempDir::new().unwrap();
        let ws = tmpdir.path().to_path_buf();
        let mut session = Session::new(
            "old-model".into(),
            ws.clone(),
            false,
            false,
            PathBuf::from("/tmp/notes"),
            PathBuf::from("/tmp/mcp"),
        );
        let mut config_workspace = PathBuf::from("/other");
        let mut config_model = "old-model".to_string();
        let messages = vec![Message {
            role: "user".to_string(),
            content: vec![ContentBlock::Text {
                text: "hello".into(),
                cache_control: None,
            }],
        }];
        apply_sync_session_payload(
            &mut session,
            &mut config_workspace,
            &mut config_model,
            messages.clone(),
            None,
            "auto".into(),
            ws.clone(),
        );
        assert_eq!(session.messages.len(), 1);
        assert_eq!(session.messages[0].role, "user");
        assert!(session.auto_model);
        assert_eq!(session.workspace, ws);
        assert_eq!(config_workspace, ws);
        assert_eq!(config_model, "auto");
    }

    fn user_msg(text: &str) -> Message {
        Message {
            role: "user".to_string(),
            content: vec![ContentBlock::Text {
                text: text.into(),
                cache_control: None,
            }],
        }
    }

    fn assistant_msg(text: &str) -> Message {
        Message {
            role: "assistant".to_string(),
            content: vec![ContentBlock::Text {
                text: text.into(),
                cache_control: None,
            }],
        }
    }

    #[test]
    fn truncate_before_last_user_message_removes_tail_exchange() {
        let mut messages = vec![
            user_msg("first"),
            assistant_msg("reply"),
            user_msg("second"),
            assistant_msg("partial"),
        ];
        assert!(truncate_before_last_user_message(&mut messages));
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[1].role, "assistant");
    }

    #[test]
    fn truncate_before_last_user_message_noop_without_user() {
        let mut messages = vec![assistant_msg("only assistant")];
        assert!(!truncate_before_last_user_message(&mut messages));
        assert_eq!(messages.len(), 1);
    }
}