huskies/server/src/crdt_state/lww_maps/llm_sessions.rs

//! Read/write helpers for the `llm_sessions` LWW-map collection, including the
//! atomic `assemble_and_advance_session` helper used by the Matrix bot.
//!
//! LLM sessions are keyed by `session_id` (typically a Matrix room ID) and track
//! per-sled high-water marks so that `assemble_and_advance_session` can inject
//! only events the LLM has not yet seen and advance the marks atomically within
//! a single CRDT lock acquisition.

use std::collections::BTreeMap;

use bft_json_crdt::json_crdt::{JsonValue, *};
use bft_json_crdt::op::ROOT_ID;
use serde_json::json;

use super::super::state::{apply_and_persist, get_crdt, rebuild_llm_session_index};
use super::super::types::{LlmSessionCrdt, LlmSessionView};
use super::event_log::GAP_PIPELINE_EVENT;

/// Write or upsert an LLM session entry keyed by `session_id`.
///
/// Creates a new entry if `session_id` is not yet present; updates
/// `persona_name` and `scope` on an existing entry.  The `high_water`
/// register is not touched by this function — use `assemble_and_advance_session`
/// to advance it atomically.
pub fn write_llm_session(session_id: &str, persona_name: &str, scope: &str) {
    let Some(state_mutex) = get_crdt() else {
        return;
    };
    let Ok(mut state) = state_mutex.lock() else {
        return;
    };

    if let Some(&idx) = state.llm_session_index.get(session_id) {
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.llm_sessions[idx]
                .persona_name
                .set(persona_name.to_string())
        });
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.llm_sessions[idx].scope.set(scope.to_string())
        });
    } else {
        let entry: JsonValue = json!({
            "session_id": session_id,
            "persona_name": persona_name,
            "scope": scope,
            "high_water": "{}",
        })
        .into();
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.llm_sessions.insert(ROOT_ID, entry)
        });
        state.llm_session_index = rebuild_llm_session_index(&state.crdt);
    }
}

/// Read a single LLM session entry by `session_id`.
pub fn read_llm_session(session_id: &str) -> Option<LlmSessionView> {
    let state_mutex = get_crdt()?;
    let state = state_mutex.lock().ok()?;
    let &idx = state.llm_session_index.get(session_id)?;
    extract_llm_session_view(&state.crdt.doc.llm_sessions[idx])
}

/// Atomically read new event-log entries for `session_id` past the stored
/// high-water marks, render them as a block of audit lines, and advance the
/// marks to prevent double-injection on the next call.
///
/// Scope is "single-sled": only events recorded by the local node (identified
/// via [`crate::crdt_state::our_node_id`]) are included.  Events from other
/// sleds are ignored in this story.
///
/// Returns an empty `Vec` when there are no new events or the CRDT is not
/// initialised.
pub fn assemble_and_advance_session(session_id: &str) -> Vec<String> {
    let local_sled_id = crate::crdt_state::our_node_id().unwrap_or_default();
    if local_sled_id.is_empty() {
        return Vec::new();
    }

    let Some(state_mutex) = get_crdt() else {
        return Vec::new();
    };
    let Ok(mut state) = state_mutex.lock() else {
        return Vec::new();
    };

    // Read the current high-water map for this session.
    let current_high_water: BTreeMap<String, u64> = {
        match state.llm_session_index.get(session_id).copied() {
            Some(idx) => parse_high_water(&state.crdt.doc.llm_sessions[idx]),
            None => BTreeMap::new(),
        }
    };

    let last_seen = current_high_water.get(&local_sled_id).copied();

    // Collect new events from the local sled past the high-water mark.
    let new_events: Vec<(u64, String, String, String, String)> = state
        .crdt
        .doc
        .event_log
        .iter()
        .filter_map(|e| extract_new_event(e, &local_sled_id, last_seen))
        .collect();

    if new_events.is_empty() {
        return Vec::new();
    }

    // Advance the high-water mark to the maximum new event_seq.
    let new_max_seq = new_events.iter().map(|(seq, ..)| *seq).max().unwrap_or(0);
    let mut new_high_water = current_high_water;
    new_high_water.insert(local_sled_id.clone(), new_max_seq);
    let new_hw_json = serde_json::to_string(&new_high_water).unwrap_or_else(|_| "{}".to_string());

    // Upsert the session entry with the new high-water value.
    let idx_opt = state.llm_session_index.get(session_id).copied();
    if let Some(idx) = idx_opt {
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.llm_sessions[idx]
                .high_water
                .set(new_hw_json.clone())
        });
    } else {
        let entry: JsonValue = json!({
            "session_id": session_id,
            "persona_name": "",
            "scope": "single-sled",
            "high_water": new_hw_json,
        })
        .into();
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.llm_sessions.insert(ROOT_ID, entry)
        });
        state.llm_session_index = rebuild_llm_session_index(&state.crdt);
    }

    // Observability: log event-log size and gap count for this sled.
    let total_entries = state
        .crdt
        .doc
        .event_log
        .iter()
        .filter(|e| matches!(e.sled_id.view(), JsonValue::String(s) if s == local_sled_id))
        .count();
    let gap_count = state
        .crdt
        .doc
        .event_log
        .iter()
        .filter(|e| {
            matches!(e.sled_id.view(), JsonValue::String(s) if s == local_sled_id)
                && matches!(e.pipeline_event.view(), JsonValue::String(s) if s == GAP_PIPELINE_EVENT)
        })
        .count();
    crate::slog!(
        "[event-log] assemble session={session_id} sled_entries={total_entries} gap_count={gap_count}"
    );

    // Render each new event as a compact audit line; gap sentinels get a
    // human-readable message so the LLM is never presented with raw field data.
    new_events
        .into_iter()
        .map(|(_, story_id, from_stage, to_stage, pipeline_event)| {
            if pipeline_event == GAP_PIPELINE_EVENT {
                format!("events between {from_stage} and {to_stage} were dropped")
            } else {
                format!(
                    "pipeline_event story_id=\"{story_id}\" from=\"{from_stage}\" \
                     to=\"{to_stage}\" event=\"{pipeline_event}\""
                )
            }
        })
        .collect()
}

/// Decode the high-water JSON string from an `LlmSessionCrdt` entry.
fn parse_high_water(entry: &LlmSessionCrdt) -> BTreeMap<String, u64> {
    match entry.high_water.view() {
        JsonValue::String(s) if !s.is_empty() && s != "{}" => {
            serde_json::from_str(&s).unwrap_or_default()
        }
        _ => BTreeMap::new(),
    }
}

/// Extract one event log entry if it belongs to `sled_id` and has an
/// `event_seq` strictly greater than `last_seen` (or `last_seen` is `None`,
/// meaning all events from this sled are new).
fn extract_new_event(
    e: &crate::crdt_state::types::EventLogEntryCrdt,
    sled_id: &str,
    last_seen: Option<u64>,
) -> Option<(u64, String, String, String, String)> {
    let entry_sled = match e.sled_id.view() {
        JsonValue::String(s) if s == sled_id => s,
        _ => return None,
    };
    let event_seq = match e.event_seq.view() {
        JsonValue::Number(n) => n as u64,
        _ => return None,
    };
    // Skip if we've already injected this event.
    if last_seen.is_some_and(|last| event_seq <= last) {
        return None;
    }
    let story_id = match e.story_id.view() {
        JsonValue::String(s) => s,
        _ => String::new(),
    };
    let from_stage = match e.from_stage.view() {
        JsonValue::String(s) => s,
        _ => String::new(),
    };
    let to_stage = match e.to_stage.view() {
        JsonValue::String(s) => s,
        _ => String::new(),
    };
    let pipeline_event = match e.pipeline_event.view() {
        JsonValue::String(s) => s,
        _ => String::new(),
    };
    let _ = entry_sled; // used only for filtering above
    Some((event_seq, story_id, from_stage, to_stage, pipeline_event))
}

/// Convert a CRDT LLM session entry into its read-only view representation.
pub(super) fn extract_llm_session_view(entry: &LlmSessionCrdt) -> Option<LlmSessionView> {
    let session_id = match entry.session_id.view() {
        JsonValue::String(s) if !s.is_empty() => s,
        _ => return None,
    };
    let persona_name = match entry.persona_name.view() {
        JsonValue::String(s) => s,
        _ => String::new(),
    };
    let scope = match entry.scope.view() {
        JsonValue::String(s) => s,
        _ => String::new(),
    };
    let high_water = parse_high_water(entry);
    Some(LlmSessionView {
        session_id,
        persona_name,
        scope,
        high_water,
    })
}