shuvarie-core 0.3.3

use std::collections::{BTreeSet, HashMap, HashSet};

use shuvarie_db::{ReasoningSegment, StoredMessage, StoredScroll, StoredSession, TextSegment};
use shuvarie_llm::ChatMsg;
use shuvarie_llm::{Role, TokenUsage};

use crate::tool_record::ToolRecord;

/// One entry of the session's per-model usage: the model code (`org/model`)
/// plus the non-Default scene that first used it (`None` when the model has
/// only served the Default scene). Deduped by code, ordered by first use;
/// drives the `/assisted-by` popup.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelUsage {
    pub code: String,
    pub scene: Option<String>,
}

/// A stored message's main-request usage: the persisted `request` payload when
/// the turn carried one, else the row's usage columns (older sessions and
/// redone turns persisted only the combined per-row totals). Rows without any
/// reported usage yield `None`.
fn request_usage_of(m: &StoredMessage) -> Option<TokenUsage> {
    if m.request.total_tokens.unwrap_or(0) > 0 {
        return Some(m.request);
    }
    let usage = TokenUsage {
        input_tokens: Some(m.input_tokens),
        output_tokens: Some(m.output_tokens),
        total_tokens: Some(m.total_tokens),
        cached_input_tokens: Some(m.cached_input_tokens),
        reasoning_tokens: Some(m.reasoning_tokens),
        ..TokenUsage::default()
    };
    (usage.total_tokens.unwrap_or(0) > 0).then_some(usage)
}

/// One tool call attached to a tree node (popup display).
#[derive(Debug, Clone)]
pub struct TreeNodeTool {
    pub name: String,
    pub ok: bool,
    pub killed: bool,
    pub worker: Option<String>,
}

/// A node of the session tree: one stored message (a user prompt, an
/// assistant reply, or a compaction summary) with its parent link and tool
/// calls. The active path runs from the leaf up through `parent` links.
#[derive(Debug, Clone)]
pub struct TreeNode {
    pub id: u64,
    pub parent: Option<u64>,
    pub role: Role,
    pub seq: u64,
    pub content: String,
    pub summary: bool,
    pub interrupted: bool,
    pub tools: Vec<TreeNodeTool>,
    /// Attachment metadata carried by the message (images, documents); the
    /// blob bytes load separately from the store.
    pub attachments: Vec<shuvarie_llm::Attachment>,
    /// The node lies on the active path (the conversation currently shown).
    pub on_path: bool,
}

#[derive(Debug, Clone, Default)]
pub struct Session {
    pub id: Option<uuid::Uuid>,
    pub title: Option<String>,
    /// The active path (root → tip) as chat messages. Forks and reloads
    /// rebuild it from the tree; live appends extend the tip.
    pub messages: Vec<ChatMsg>,
    pub reasoning: HashMap<u64, Vec<ReasoningSegment>>,
    /// The turn's text runs keyed by dense message index, with the tool-call
    /// positions they streamed at; drives the reload interleave.
    pub text_segments: HashMap<u64, Vec<TextSegment>>,
    pub interrupted: HashMap<u64, bool>,
    /// Tool calls belonging to the active path, keyed by dense message index.
    pub tool_records: Vec<ToolRecord>,
    pub tokens: u64,
    pub cost: f64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub reasoning_tokens: u64,
    pub cached_tokens: u64,
    /// Dense indices (within `messages`) of the summary nodes on the active
    /// path. `history_for_send` drops everything before the newest one.
    pub summaries: BTreeSet<u64>,
    /// The full stored tree (all branches, not just the active path); drives
    /// the `/tree` popup. Only populated on load.
    pub nodes: Vec<TreeNode>,
    /// DB id of the active path's tip: the parent new appends hang from.
    pub leaf_id: Option<u64>,
    /// The session's active scene; `None` = the built-in default scene.
    pub scene: Option<String>,
    /// The scene the model was last told about — what the DB `sessions.scene`
    /// column holds. Differs from `scene` only inside the deferral window
    /// after a mid-session switch: the switch moves `scene` at once, while
    /// the announce (interlude injection + DB write) waits for the first
    /// request under it. Seeded from storage on load and kept in step by the
    /// announce.
    pub announced_scene: Option<String>,
    /// Usage of the most recent main-stream request against the current
    /// post-compaction conversation. Maintained live by the turn loop (every
    /// main-stream usage report) and seeded from storage on load — only from
    /// requests made after the newest compaction summary, since earlier usage
    /// describes the conversation the summary replaced. The next turn seeds
    /// the LLM's measured compaction trigger and the sidebar's context anchor
    /// from it.
    pub last_usage: Option<TokenUsage>,
    /// Per-model usage on the active path, deduped by model code and ordered
    /// by first use; `scene` carries the first non-Default scene the model
    /// served. Rebuilt from the stored rows on load; live turns extend it.
    pub models_used: Vec<ModelUsage>,
    /// Chat pane scroll position persisted at the last leave of the session;
    /// the TUI seeds its scroll engine from it when the session loads.
    pub scroll: StoredScroll,
}

/// The request that resumes the active path's last turn after a failed
/// attempt: a turn retry (a connection drop, a worker failure) or an overflow
/// auto-continue.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TurnResume {
    /// The attempt produced no reply to continue from: the turn's stored
    /// prompt is replayed as the request's prompt, with the
    /// summary-truncated history before it. The prompt's node is not part of
    /// `prior`, so re-sending its attachments is the caller's to arrange.
    Replay { prompt: String, prior: Vec<ChatMsg> },
    /// The attempt left a reply row on the active path after the turn's
    /// prompt: everything before the tip — the prompt node included — rides
    /// the request's history, so the model continues the interrupted turn
    /// instead of replaying it from its prompt. A replay redid the work the
    /// attempt had already finished and, for a session whose turn is its
    /// first prompt (the common single-task session), re-sent the whole task
    /// with no history at all — the model read it as a fresh request and
    /// started the task over. A tip that wrote text is the request's last
    /// history message; a tip that wrote none (a reasoning- or tool-only
    /// attempt) is dropped, since a blank assistant message carries nothing.
    /// The caller prompts with the resume cue.
    Continue { prior: Vec<ChatMsg> },
}

impl Session {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn from_stored(stored: StoredSession) -> Self {
        let mut nodes: Vec<TreeNode> = stored
            .messages
            .iter()
            .map(|m| TreeNode {
                id: m.id,
                parent: m.parent_id,
                role: m.role.into(),
                seq: m.seq,
                content: m.content.clone(),
                summary: m.summary,
                interrupted: m.interrupted,
                tools: Vec::new(),
                attachments: m.attachments.clone(),
                on_path: false,
            })
            .collect();
        let mut tools_by_msg: HashMap<u64, Vec<TreeNodeTool>> = HashMap::new();
        for tc in &stored.tool_calls {
            tools_by_msg
                .entry(tc.message_id)
                .or_default()
                .push(TreeNodeTool {
                    name: tc.name.clone(),
                    ok: tc.ok,
                    killed: tc.killed,
                    worker: tc.worker.clone(),
                });
        }
        for node in &mut nodes {
            if let Some(tools) = tools_by_msg.remove(&node.id) {
                node.tools = tools;
            }
        }
        let by_id: HashMap<u64, usize> = nodes.iter().enumerate().map(|(i, n)| (n.id, i)).collect();

        // The active tip: `Some(EMPTY_LEAF)` marks a cleared (empty) path; a
        // stored leaf that still resolves wins; otherwise — legacy unset or
        // dangling — the newest message by seq.
        let newest = || {
            nodes
                .iter()
                .enumerate()
                .max_by_key(|(_, n)| n.seq)
                .map(|(i, _)| i)
        };
        let leaf_pos = match stored.leaf_id {
            Some(shuvarie_db::EMPTY_LEAF) => None,
            Some(id) => by_id.get(&id).copied().or_else(newest),
            None => newest(),
        };

        // Walk leaf → root, reversing into path order. The `seen` guard makes
        // a corrupt parent cycle terminate instead of looping.
        let mut path_ids: Vec<u64> = Vec::with_capacity(nodes.len());
        let mut seen: HashSet<u64> = HashSet::new();
        if let Some(pos) = leaf_pos {
            let mut cur = Some(nodes[pos].id);
            while let Some(id) = cur {
                if !seen.insert(id) {
                    break;
                }
                path_ids.push(id);
                cur = by_id.get(&id).copied().and_then(|p| nodes[p].parent);
            }
            path_ids.reverse();
            for id in &path_ids {
                if let Some(&pos) = by_id.get(id) {
                    nodes[pos].on_path = true;
                }
            }
        }
        let dense_of: HashMap<u64, usize> = path_ids
            .iter()
            .enumerate()
            .map(|(i, id)| (*id, i))
            .collect();

        let mut s = Self {
            id: Some(stored.id),
            title: Some(stored.title),
            leaf_id: path_ids.last().copied(),
            announced_scene: stored.scene.clone(),
            scene: stored.scene,
            scroll: stored.scroll,
            ..Self::default()
        };
        // Usage anchoring: requests made after the newest compaction summary
        // describe the current conversation's size. Everything before it was
        // replaced by that summary, so its stored usage would over-state the
        // context — and re-arm the measured compaction trigger against a
        // history that just shrank. With no summary on the path, every
        // request qualifies.
        let newest_summary: Option<u64> = path_ids
            .iter()
            .enumerate()
            .filter(|(_, id)| stored.messages[by_id[id]].summary)
            .map(|(idx, _)| idx as u64)
            .next_back();
        for (idx, id) in path_ids.iter().enumerate() {
            let m = &stored.messages[by_id[id]];
            s.input_tokens = s.input_tokens.saturating_add(m.input_tokens);
            s.output_tokens = s.output_tokens.saturating_add(m.output_tokens);
            s.tokens = s.tokens.saturating_add(m.total_tokens);
            s.reasoning_tokens = s.reasoning_tokens.saturating_add(m.reasoning_tokens);
            s.cached_tokens = s.cached_tokens.saturating_add(m.cached_input_tokens);
            s.cost += m.cost;
            s.messages.push(ChatMsg {
                role: m.role.into(),
                content: m.content.clone(),
                attachments: m.attachments.clone(),
            });
            let idx = idx as u64;
            if !m.reasoning.is_empty() {
                s.reasoning.insert(idx, m.reasoning.clone());
            }
            if !m.text_segments.is_empty() {
                s.text_segments.insert(idx, m.text_segments.clone());
            }
            if m.interrupted {
                s.interrupted.insert(idx, true);
            }
            if m.summary {
                s.summaries.insert(idx);
            } else if newest_summary.is_none_or(|pos| idx > pos)
                && let Some(usage) = request_usage_of(m)
            {
                s.last_usage = Some(usage);
            }
            if !m.summary {
                s.record_model_use(m.model_code.as_deref(), m.scene.as_deref());
            }
        }
        s.tool_records = stored
            .tool_calls
            .iter()
            .filter_map(|tc| {
                let idx = dense_of.get(&tc.message_id).copied()?;
                let mut record = ToolRecord::from_stored(tc.clone());
                record.message_seq = idx as u64;
                Some(record)
            })
            .collect();
        s.nodes = nodes;
        s
    }

    pub fn clear(&mut self) {
        self.id = None;
        self.title = None;
        self.messages.clear();
        self.reasoning.clear();
        self.text_segments.clear();
        self.interrupted.clear();
        self.tool_records.clear();
        self.tokens = 0;
        self.cost = 0.0;
        self.input_tokens = 0;
        self.output_tokens = 0;
        self.reasoning_tokens = 0;
        self.cached_tokens = 0;
        self.summaries.clear();
        self.nodes.clear();
        self.leaf_id = None;
        self.scene = None;
        self.announced_scene = None;
        self.last_usage = None;
        self.models_used.clear();
        self.scroll = StoredScroll::default();
    }

    pub fn push_user(&mut self, content: impl Into<String>) {
        self.messages.push(ChatMsg::user(content));
    }

    /// `push_user` carrying attachment metadata (resolved and prepared by
    /// the turn's resolve step before anything persists).
    pub fn push_user_with(
        &mut self,
        content: impl Into<String>,
        attachments: Vec<shuvarie_llm::Attachment>,
    ) {
        let mut msg = ChatMsg::user(content);
        msg.attachments = attachments;
        self.messages.push(msg);
    }

    pub fn push_assistant(&mut self, content: impl Into<String>) {
        self.messages.push(ChatMsg::assistant(content));
    }

    /// Record one model use against the session: deduped by model code,
    /// ordered by first use, with the first non-Default scene kept as the
    /// annotation. A `None` code (the model is not in the catalog) is a no-op.
    pub fn record_model_use(&mut self, model_code: Option<&str>, scene: Option<&str>) {
        let Some(code) = model_code else {
            return;
        };
        let scene = scene.filter(|scene| !scene.is_empty()).map(str::to_string);
        match self.models_used.iter_mut().find(|u| u.code == code) {
            Some(entry) => {
                if entry.scene.is_none() {
                    entry.scene = scene;
                }
            }
            None => self.models_used.push(ModelUsage {
                code: code.to_string(),
                scene,
            }),
        }
    }

    /// The active path's last user prompt node, if any.
    pub fn last_user_node(&self) -> Option<&TreeNode> {
        self.nodes
            .iter()
            .rev()
            .find(|n| n.on_path && n.role == Role::User)
    }

    /// The request that resumes a failed turn: the attempt's own output must
    /// reach the retried request — leaving it out is what makes a retried
    /// turn replay from its prompt and redo work the attempt had already
    /// finished.
    ///
    /// An assistant tip is the failed attempt's own reply row — the model's
    /// progress log for the turn — and empty when the attempt reasoned or
    /// called tools but never wrote text. Either way the request continues
    /// the turn ([`TurnResume::Continue`]): the whole active path —
    /// summary-truncated, prompt node included, an empty tip dropped — rides
    /// the history, so the model picks the work already on the branch back up
    /// instead of reading the turn's prompt as a fresh request. Only a path
    /// whose tip is the user prompt itself — the attempt left no reply row at
    /// all — replays that prompt ([`TurnResume::Replay`]) with the history
    /// before it ([`Self::history_for_send_until`]). `None` when the path has
    /// no user prompt to resume from.
    pub fn last_turn_resume(&self) -> Option<TurnResume> {
        if let Some(tip) = self.messages.last()
            && tip.role == Role::Assistant
        {
            let end = if tip.content.is_empty() {
                self.messages.len() - 1
            } else {
                self.messages.len()
            };
            let prior = self.history_for_send_until(end);
            // A lone empty assistant row has no turn behind it to resume.
            if prior.iter().any(|m| m.role == Role::User) {
                return Some(TurnResume::Continue { prior });
            }
        }
        let end = self.messages.iter().rposition(|m| m.role == Role::User)?;
        Some(TurnResume::Replay {
            prompt: self.messages[end].content.clone(),
            prior: self.history_for_send_until(end),
        })
    }

    /// The active path's first user prompt (root → tip): the source for LLM
    /// title drafting. Compaction summaries (assistant role) never match.
    pub fn first_user_prompt(&self) -> Option<&str> {
        self.messages
            .iter()
            .find(|m| m.role == Role::User)
            .map(|m| m.content.as_str())
    }

    pub fn last_assistant_interrupted(&self) -> bool {
        let Some(idx) = self
            .messages
            .iter()
            .rposition(|m| m.role == shuvarie_llm::Role::Assistant)
        else {
            return false;
        };
        self.interrupted
            .get(&(idx as u64))
            .copied()
            .unwrap_or(false)
    }

    pub fn add_usage(&mut self, usage: TokenUsage, cost: f64) {
        self.tokens = self.tokens.saturating_add(usage.total_tokens.unwrap_or(0));
        self.cost += cost;
        self.input_tokens = self
            .input_tokens
            .saturating_add(usage.input_tokens.unwrap_or(0));
        self.output_tokens = self
            .output_tokens
            .saturating_add(usage.output_tokens.unwrap_or(0));
        self.reasoning_tokens = self
            .reasoning_tokens
            .saturating_add(usage.reasoning_tokens.unwrap_or(0));
        self.cached_tokens = self
            .cached_tokens
            .saturating_add(usage.cached_input_tokens.unwrap_or(0));
    }

    /// The session's cumulative usage as a [`TokenUsage`], for usage snapshots.
    pub fn usage(&self) -> TokenUsage {
        TokenUsage {
            total_tokens: Some(self.tokens),
            input_tokens: Some(self.input_tokens),
            output_tokens: Some(self.output_tokens),
            cached_input_tokens: Some(self.cached_tokens),
            reasoning_tokens: Some(self.reasoning_tokens),
            ..TokenUsage::default()
        }
    }

    /// Build the chat history to send to the LLM for a new turn: everything
    /// before the pending user message (the active path's tip), truncated at
    /// the newest summary node — the summary replaces the history before it.
    pub fn history_for_send(&self) -> Vec<ChatMsg> {
        self.history_for_send_until(self.messages.len().saturating_sub(1))
    }

    /// [`Self::history_for_send`] over `messages[..end]` — the history
    /// preceding message index `end`, truncated at the newest summary inside
    /// the span. A turn resume passes the index of the turn's pending prompt
    /// so the partial assistant tip after it stays out of the retried
    /// request.
    pub fn history_for_send_until(&self, end: usize) -> Vec<ChatMsg> {
        if end == 0 {
            return Vec::new();
        }
        let begin = self
            .summaries
            .iter()
            .rev()
            .find(|&&s| (s as usize) < end)
            .map(|&s| s as usize)
            .unwrap_or(0);
        self.messages[begin..end].to_vec()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use shuvarie_db::MsgRole;

    fn stored_message(seq: u64, role: MsgRole, content: &str) -> shuvarie_db::StoredMessage {
        shuvarie_db::StoredMessage {
            id: seq + 1,
            parent_id: None,
            role,
            content: content.to_string(),
            reasoning: Vec::new(),
            text_segments: Vec::new(),
            interrupted: false,
            seq,
            input_tokens: 0,
            output_tokens: 0,
            total_tokens: 0,
            cached_input_tokens: 0,
            reasoning_tokens: 0,
            cost: 0.0,
            summary: false,
            request: TokenUsage::default(),
            model_code: None,
            scene: None,
            attachments: Vec::new(),
        }
    }

    fn stored_session(messages: Vec<shuvarie_db::StoredMessage>) -> shuvarie_db::StoredSession {
        shuvarie_db::StoredSession {
            scene: None,
            id: uuid::Uuid::now_v7(),
            title: "t".into(),
            provider: None,
            model: None,
            leaf_id: None,
            messages,
            tool_calls: Vec::new(),
            scroll: shuvarie_db::StoredScroll::default(),
            created_at: jiff::Timestamp::now(),
            updated_at: jiff::Timestamp::now(),
        }
    }

    #[test]
    fn from_stored_seeds_last_usage_from_the_request_payload() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "hi"),
            stored_message(1, MsgRole::Assistant, "ok"),
        ]);
        stored.messages[1].input_tokens = 12_000;
        stored.messages[1].output_tokens = 200;
        stored.messages[1].total_tokens = 12_200;
        stored.messages[1].cached_input_tokens = 11_000;
        stored.messages[1].request = TokenUsage {
            input_tokens: Some(20_000),
            output_tokens: Some(200),
            total_tokens: Some(20_200),
            cached_input_tokens: Some(19_400),
            ..TokenUsage::default()
        };

        let session = Session::from_stored(stored);
        let last = session.last_usage.expect("request usage restored");
        assert_eq!(last.total_tokens, Some(20_200));
        assert_eq!(last.cached_input_tokens, Some(19_400));
    }

    #[test]
    fn from_stored_falls_back_to_row_usage_without_a_request_payload() {
        let mut stored = stored_session(vec![stored_message(0, MsgRole::Assistant, "ok")]);
        stored.messages[0].input_tokens = 500;
        stored.messages[0].output_tokens = 100;
        stored.messages[0].total_tokens = 600;

        let session = Session::from_stored(stored);
        let last = session.last_usage.expect("row usage restored");
        assert_eq!(last.total_tokens, Some(600));
    }

    #[test]
    fn from_stored_skips_summary_rows_and_pre_summary_usage_for_last_usage() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::Assistant, "ok"),
            stored_message(1, MsgRole::User, "summary"),
            stored_message(2, MsgRole::User, "again"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[2].id);
        stored.messages[0].total_tokens = 30;
        stored.messages[0].request = TokenUsage {
            total_tokens: Some(30),
            ..TokenUsage::default()
        };
        stored.messages[1].summary = true;
        stored.messages[1].total_tokens = 9_999;
        stored.messages[1].request = TokenUsage {
            total_tokens: Some(9_999),
            ..TokenUsage::default()
        };

        let session = Session::from_stored(stored);
        // Both usable rows predate the newest summary: the head they measured
        // was replaced by that summary, so nothing on the path describes the
        // current conversation and the anchor stays unset.
        assert!(session.last_usage.is_none());
    }

    #[test]
    fn from_stored_seeds_last_usage_from_post_summary_requests() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::Assistant, "ok"),
            stored_message(1, MsgRole::User, "summary"),
            stored_message(2, MsgRole::Assistant, "fresh"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[2].id);
        stored.messages[0].request = TokenUsage {
            total_tokens: Some(99_000),
            ..TokenUsage::default()
        };
        stored.messages[1].summary = true;
        stored.messages[2].request = TokenUsage {
            total_tokens: Some(3_000),
            ..TokenUsage::default()
        };

        let session = Session::from_stored(stored);
        let last = session.last_usage.expect("post-summary request restored");
        assert_eq!(
            last.total_tokens,
            Some(3_000),
            "pre-summary usage is ignored"
        );
    }

    /// Post-`/compact` layout: the summary is spliced mid-path and the kept
    /// tail hangs under it, but the tail's requests chronologically predate
    /// the summary (they measured the pre-compaction context).
    #[test]
    fn from_stored_seeds_last_usage_from_the_kept_compaction_tail() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "one"),
            stored_message(1, MsgRole::Assistant, "r1"),
            stored_message(2, MsgRole::User, "summary"),
            stored_message(3, MsgRole::User, "two"),
            stored_message(4, MsgRole::Assistant, "r2"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[4].id);
        stored.messages[1].request = TokenUsage {
            total_tokens: Some(30_000),
            ..TokenUsage::default()
        };
        stored.messages[2].summary = true;
        stored.messages[4].request = TokenUsage {
            total_tokens: Some(32_000),
            ..TokenUsage::default()
        };

        let session = Session::from_stored(stored);
        let last = session.last_usage.expect("kept-tail request restored");
        assert_eq!(last.total_tokens, Some(32_000));
    }

    #[test]
    fn from_stored_without_reported_usage_leaves_last_usage_none() {
        let stored = stored_session(vec![stored_message(0, MsgRole::User, "hi")]);
        let session = Session::from_stored(stored);
        assert!(session.last_usage.is_none());
    }

    #[test]
    fn from_stored_maps_text_segments_to_dense_indices() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "go"),
            stored_message(7, MsgRole::Assistant, "start\n\nresumed"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[1].id);
        stored.messages[1].text_segments = vec![
            shuvarie_db::TextSegment {
                after_tool: 0,
                text: "start".into(),
            },
            shuvarie_db::TextSegment {
                after_tool: 1,
                text: "\n\nresumed".into(),
            },
        ];

        let session = Session::from_stored(stored);
        let runs = session
            .text_segments
            .get(&1)
            .expect("segments keyed by dense index");
        assert_eq!(runs.len(), 2);
        assert_eq!(runs[1].after_tool, 1);
    }

    /// Chain helper: link each message to the previous one via `parent_id`.
    fn chain(messages: &mut [shuvarie_db::StoredMessage]) {
        for i in 1..messages.len() {
            messages[i].parent_id = Some(messages[i - 1].id);
        }
    }

    #[test]
    fn from_stored_walks_the_active_path_from_the_leaf() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "one"),
            stored_message(1, MsgRole::Assistant, "r1"),
            stored_message(2, MsgRole::User, "two"),
            stored_message(3, MsgRole::Assistant, "r2"),
            stored_message(3, MsgRole::User, "fork prompt"),
            stored_message(4, MsgRole::Assistant, "fork reply"),
        ]);
        chain(&mut stored.messages);
        // Fork: the last two messages branch off the first assistant.
        stored.messages[4].parent_id = Some(stored.messages[1].id);
        stored.messages[5].parent_id = Some(stored.messages[4].id);
        stored.leaf_id = Some(stored.messages[5].id);

        let session = Session::from_stored(stored);
        let path: Vec<&str> = session
            .messages
            .iter()
            .map(|m| m.content.as_str())
            .collect();
        assert_eq!(path, vec!["one", "r1", "fork prompt", "fork reply"]);
        assert_eq!(session.nodes.iter().filter(|n| n.on_path).count(), 4);
        let on_fork = session.nodes.iter().find(|n| n.content == "r2").unwrap();
        assert!(!on_fork.on_path);
    }

    #[test]
    fn from_stored_falls_back_to_newest_message_without_leaf() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "one"),
            stored_message(1, MsgRole::Assistant, "r1"),
        ]);
        chain(&mut stored.messages);

        let session = Session::from_stored(stored);
        assert_eq!(session.messages.len(), 2);
        assert_eq!(session.messages[1].content, "r1");
    }

    #[test]
    fn from_stored_maps_attachments_to_the_path_messages_and_nodes() {
        let mut stored = stored_session(vec![stored_message(0, MsgRole::User, "see")]);
        stored.messages[0].attachments = vec![shuvarie_llm::Attachment {
            kind: shuvarie_llm::AttachmentKind::Image,
            name: "photo.png".into(),
            media_type: "image/png".into(),
            size: 2_000,
            sha256: "a2".repeat(32),
        }];

        let session = Session::from_stored(stored);
        assert_eq!(session.messages[0].attachments.len(), 1);
        assert_eq!(session.messages[0].attachments[0].name, "photo.png");
        assert_eq!(session.nodes[0].attachments.len(), 1);
        assert_eq!(session.nodes[0].attachments[0].sha256, "a2".repeat(32));
        // The node's copy is independent of the message's (no shared mutation).
        assert_ne!(
            session.nodes[0].attachments.as_ptr(),
            session.messages[0].attachments.as_ptr()
        );
    }

    #[test]
    fn from_stored_treats_a_cleared_leaf_as_an_empty_path() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "one"),
            stored_message(1, MsgRole::Assistant, "r1"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(shuvarie_db::EMPTY_LEAF);

        let session = Session::from_stored(stored);
        assert!(
            session.messages.is_empty(),
            "undoing the first prompt clears the chat"
        );
        assert!(session.nodes.iter().all(|n| !n.on_path));
        assert_eq!(session.leaf_id, None, "new appends hang from the root");
    }

    #[test]
    fn from_stored_uses_the_walked_tip_as_the_leaf() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "one"),
            stored_message(1, MsgRole::Assistant, "r1"),
        ]);
        chain(&mut stored.messages);

        let session = Session::from_stored(stored);
        assert_eq!(
            session.leaf_id,
            Some(2),
            "the fallback tip becomes the leaf"
        );
    }

    #[test]
    fn from_stored_excludes_forked_away_usage_and_tools() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "one"),
            stored_message(1, MsgRole::Assistant, "r1"),
            stored_message(2, MsgRole::User, "fork"),
        ]);
        chain(&mut stored.messages);
        stored.messages[1].input_tokens = 100;
        stored.messages[1].total_tokens = 100;
        stored.messages[2].parent_id = Some(stored.messages[0].id);
        stored.leaf_id = Some(stored.messages[0].id);
        let mut forked_tool = tool_call_of(stored.messages[2].id);
        forked_tool.name = "read_file".into();
        stored.tool_calls = vec![forked_tool];

        let session = Session::from_stored(stored);
        assert_eq!(session.tokens, 0, "forked-away branch usage excluded");
        assert!(session.tool_records.is_empty());
    }

    fn tool_call_of(message_id: u64) -> shuvarie_db::StoredToolCall {
        shuvarie_db::StoredToolCall {
            id: 9,
            message_id,
            session_id: uuid::Uuid::nil(),
            seq: 0,
            name: "read_file".into(),
            args_json: "{}".into(),
            output: String::new(),
            stderr: String::new(),
            ok: true,
            killed: false,
            worker: None,
            file_change_json: String::new(),
            original_content: None,
            new_content: None,
            duration_ms: 0,
        }
    }

    #[test]
    fn history_for_send_stops_at_the_newest_summary() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::User, "u2"),
            stored_message(3, MsgRole::Assistant, "summary"),
            stored_message(4, MsgRole::User, "u3"),
            stored_message(5, MsgRole::Assistant, "a3"),
        ]);
        chain(&mut stored.messages);
        stored.messages[3].summary = true;
        stored.leaf_id = Some(stored.messages[5].id);

        let session = Session::from_stored(stored);
        let history: Vec<String> = session
            .history_for_send()
            .iter()
            .map(|m| m.content.clone())
            .collect();
        assert_eq!(history, vec!["summary", "u3"], "summary + tail, no tip");
    }

    #[test]
    fn history_for_send_without_summary_is_full_but_pending() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::User, "pending"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[2].id);

        let session = Session::from_stored(stored);
        let history: Vec<String> = session
            .history_for_send()
            .iter()
            .map(|m| m.content.clone())
            .collect();
        assert_eq!(history, vec!["u1", "a1"]);
        assert!(session.messages[2].content == "pending");
    }

    #[test]
    fn history_for_send_until_takes_the_span_end() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::User, "u2"),
            stored_message(3, MsgRole::Assistant, "a2 partial"),
        ]);
        chain(&mut stored.messages);
        stored.messages[3].interrupted = true;
        stored.leaf_id = Some(stored.messages[3].id);

        let session = Session::from_stored(stored);
        let until_two: Vec<String> = session
            .history_for_send_until(2)
            .iter()
            .map(|m| m.content.clone())
            .collect();
        assert_eq!(
            until_two,
            vec!["u1", "a1"],
            "history up to the turn's prompt, partial tip excluded"
        );
        assert!(session.history_for_send_until(0).is_empty());
    }

    #[test]
    fn last_turn_resume_continues_the_interrupted_reply() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::User, "u2"),
            stored_message(3, MsgRole::Assistant, "a2 partial"),
        ]);
        chain(&mut stored.messages);
        stored.messages[3].interrupted = true;
        stored.leaf_id = Some(stored.messages[3].id);

        let session = Session::from_stored(stored);
        let Some(TurnResume::Continue { prior }) = session.last_turn_resume() else {
            panic!("the partial tip is continued, not replayed");
        };
        let prior: Vec<String> = prior.iter().map(|m| m.content.clone()).collect();
        assert_eq!(
            prior,
            vec!["u1", "a1", "u2", "a2 partial"],
            "the whole path rides the request, the interrupted reply included"
        );
    }

    #[test]
    fn last_turn_resume_truncates_the_continuation_at_the_newest_summary() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::User, "u2"),
            stored_message(3, MsgRole::Assistant, "summary"),
            stored_message(4, MsgRole::User, "u3"),
            stored_message(5, MsgRole::Assistant, "a3 partial"),
        ]);
        chain(&mut stored.messages);
        stored.messages[3].summary = true;
        stored.messages[5].interrupted = true;
        stored.leaf_id = Some(stored.messages[5].id);

        let session = Session::from_stored(stored);
        let Some(TurnResume::Continue { prior }) = session.last_turn_resume() else {
            panic!("the partial tip is continued, not replayed");
        };
        let prior: Vec<String> = prior.iter().map(|m| m.content.clone()).collect();
        assert_eq!(
            prior,
            vec!["summary", "u3", "a3 partial"],
            "the summary replaces the history before it; the turn's prompt stays"
        );
    }

    #[test]
    fn last_turn_resume_replays_the_prompt_when_nothing_streamed() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::User, "u2"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[2].id);

        let session = Session::from_stored(stored);
        let Some(TurnResume::Replay { prompt, prior }) = session.last_turn_resume() else {
            panic!("a prompt with no reply replays");
        };
        assert_eq!(prompt, "u2");
        let prior: Vec<String> = prior.iter().map(|m| m.content.clone()).collect();
        assert_eq!(prior, vec!["u1", "a1"]);
    }

    #[test]
    fn last_turn_resume_continues_after_an_empty_reply() {
        // A turn cut off after its tool batch but before any text leaves an
        // empty assistant tip. There is no reply of its own to continue, but
        // the tip must not push the retry down the replay path: for a turn
        // that is the session's first prompt, replaying re-sends the task with
        // no history and the model starts the task over. The empty row is
        // dropped from the request and the path before it carries the turn.
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, ""),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[1].id);

        let session = Session::from_stored(stored);
        let Some(TurnResume::Continue { prior }) = session.last_turn_resume() else {
            panic!("an empty reply still continues the turn");
        };
        let prior: Vec<String> = prior.iter().map(|m| m.content.clone()).collect();
        assert_eq!(
            prior,
            vec!["u1"],
            "the empty tip is dropped, the turn stays"
        );
    }

    #[test]
    fn last_turn_resume_keeps_the_earlier_partial_after_an_empty_retry() {
        // The reported failure: the retried attempt reasoned or called tools
        // and dropped before writing text, so the branch tip is an empty
        // assistant row. The retry that follows must still carry the first
        // attempt's partial reply — not fall back to replaying the turn's
        // prompt, which (for the session's only prompt) would re-send the bare
        // task and start it over.
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "do the task"),
            stored_message(1, MsgRole::Assistant, "halfway through the task"),
            stored_message(2, MsgRole::Assistant, ""),
        ]);
        chain(&mut stored.messages);
        stored.messages[1].interrupted = true;
        stored.messages[2].interrupted = true;
        stored.leaf_id = Some(stored.messages[2].id);

        let session = Session::from_stored(stored);
        let Some(TurnResume::Continue { prior }) = session.last_turn_resume() else {
            panic!("an empty retry tip must not replay the task prompt");
        };
        let prior: Vec<String> = prior.iter().map(|m| m.content.clone()).collect();
        assert_eq!(
            prior,
            vec!["do the task", "halfway through the task"],
            "the earlier partial survives on the request's history"
        );
    }

    #[test]
    fn last_turn_resume_without_a_prompt_or_a_reply_is_none() {
        let mut stored = stored_session(vec![stored_message(0, MsgRole::Assistant, "")]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[0].id);

        let session = Session::from_stored(stored);
        assert!(session.last_turn_resume().is_none());
    }

    #[test]
    fn last_user_node_finds_the_active_prompt() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::User, "u2"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[2].id);

        let session = Session::from_stored(stored);
        let node = session.last_user_node().expect("user node on path");
        assert_eq!(node.content, "u2");
        assert_eq!(node.parent, Some(2), "the parent is the assistant row");
    }

    #[test]
    fn record_model_use_dedupes_by_code_and_keeps_first_non_default_scene() {
        let mut session = Session::new();
        session.record_model_use(Some("zai-org/glm-5.3-flash"), None);
        session.record_model_use(Some("deepseek-ai/deepseek-v4"), Some("Plan"));
        session.record_model_use(Some("zai-org/glm-5.3-flash"), Some("Review"));
        session.record_model_use(Some("deepseek-ai/deepseek-v4"), None);
        session.record_model_use(None, Some("Plan"));

        assert_eq!(
            session.models_used,
            vec![
                ModelUsage {
                    code: "zai-org/glm-5.3-flash".into(),
                    scene: Some("Review".into()),
                },
                ModelUsage {
                    code: "deepseek-ai/deepseek-v4".into(),
                    scene: Some("Plan".into()),
                },
            ]
        );
    }

    #[test]
    fn from_stored_rebuilds_models_used_from_the_active_path() {
        let mut stored = stored_session(vec![
            stored_message(0, MsgRole::User, "u1"),
            stored_message(1, MsgRole::Assistant, "a1"),
            stored_message(2, MsgRole::Assistant, "summary"),
            stored_message(3, MsgRole::User, "u2"),
            stored_message(4, MsgRole::Assistant, "a2"),
            stored_message(5, MsgRole::Assistant, "abandoned"),
        ]);
        chain(&mut stored.messages);
        stored.leaf_id = Some(stored.messages[4].id);
        stored.messages[1].model_code = Some("zai-org/glm-5.3-flash".into());
        stored.messages[1].scene = None;
        stored.messages[2].summary = true;
        stored.messages[2].model_code = Some("anthropic/claude-summarizer".into());
        stored.messages[4].model_code = Some("deepseek-ai/deepseek-v4".into());
        stored.messages[4].scene = Some("Plan".into());
        stored.messages[5].model_code = Some("off-path/model".into());

        let session = Session::from_stored(stored);
        assert_eq!(
            session.models_used,
            vec![
                ModelUsage {
                    code: "zai-org/glm-5.3-flash".into(),
                    scene: None,
                },
                ModelUsage {
                    code: "deepseek-ai/deepseek-v4".into(),
                    scene: Some("Plan".into()),
                },
            ]
        );
    }
}