lash-protocol-rlm 0.1.0-alpha.66

//! RLM history rendering: turns the chronological turn view into the LLM
//! message sequence the model sees each iteration.
//!
//! Contract:
//! - **Message roles.** Each prior step is a chat message, not one packed
//!   blob. User messages render as `User`; prior lashlang steps and protocol events
//!   render as `Assistant` (see `borrowed_history_entry_llm_role`).
//! - **Cache fence.** The last history message is marked with a
//!   `cache_breakpoint` (`mark_last_history_text_cache_breakpoint`) so the
//!   provider can reuse the stable history prefix across iterations. Only the
//!   volatile `=== CURRENT ITERATION ===` tail — iteration number, turn
//!   events, finalization, required-output schema, context budget — is appended
//!   uncached by `append_current_iteration_message`.
//! - **Re-fetch handle.** Outputs lossy-projected for the history prompt are
//!   tagged with `full: history[N].output[M]`. That handle resolves: the
//!   `history` projected binding carries the full untruncated value, so the
//!   model can recover it by re-printing the reference. Proven by
//!   `projection::context::tests::history_step_output_resolves_full_untruncated_value`.
//! - **Variables.** The live variable namespace is surfaced separately as the
//!   "Bound Variables" prompt contribution (`crate::rlm_support::render_bound_variables`),
//!   not here — history carries the trajectory, not current globals.

use std::collections::HashSet;
use std::fmt::Write as _;
use std::sync::Arc;

use lash_core::llm::types::{LlmAttachment, LlmContentBlock, LlmMessage, LlmRole};
use lash_core::{BorrowedChronologicalEntry, BorrowedChronologicalPayload, head_tail_truncate};
#[cfg(test)]
use lash_core::{ChronologicalEntry, ChronologicalPayload};
#[cfg(test)]
use lash_rlm_types::RlmHistoryItem;
use lash_rlm_types::{RlmAttachmentRef, RlmHistoryRole, RlmImageRef};
use lashlang::{Value as FlowValue, ValueProjectionContext};

use crate::projection::{decode_rlm_protocol_event, json_to_flow_value};

pub(super) struct RlmHistoryRenderInput<'a> {
    pub(super) events: &'a [lash_core::SessionEventRecord],
    pub(super) turn_messages: &'a lash_core::MessageSequence,
    pub(super) turn_causes: &'a [lash_core::TurnCause],
    pub(super) max_output_chars: usize,
    pub(super) protocol_iteration: usize,
    pub(super) finalization: &'a str,
    pub(super) required_output: Option<&'a str>,
    pub(super) final_answer_format: Option<&'a str>,
    pub(super) budget_suffix: Option<&'a str>,
}

#[derive(Clone, Copy)]
pub(super) struct CurrentIterationMessageInput<'a> {
    pub(super) saw_history: bool,
    pub(super) protocol_iteration: usize,
    pub(super) turn_causes: &'a [lash_core::TurnCause],
    pub(super) finalization: &'a str,
    pub(super) required_output: Option<&'a str>,
    pub(super) final_answer_format: Option<&'a str>,
    pub(super) budget_suffix: Option<&'a str>,
}

#[cfg(test)]
pub(super) struct RlmHistoryTestRenderInput<'a> {
    pub(super) projection: &'a lash_core::ChronologicalProjection,
    pub(super) max_output_chars: usize,
    pub(super) protocol_iteration: usize,
    pub(super) finalization: &'a str,
    pub(super) required_output: Option<&'a str>,
    pub(super) final_answer_format: Option<&'a str>,
    pub(super) budget_suffix: Option<&'a str>,
}

pub(super) fn build_rlm_history_messages_from_turn(
    input: RlmHistoryRenderInput<'_>,
    attachments: &mut Vec<LlmAttachment>,
) -> Vec<LlmMessage> {
    let mut messages = Vec::new();
    let mut saw_history = false;
    let active_cause_ids = input
        .turn_causes
        .iter()
        .map(|cause| cause.id.as_str())
        .collect::<HashSet<_>>();
    lash_core::visit_turn_view(input.events, input.turn_messages, |entry| {
        if borrowed_entry_is_active_cause(entry, &active_cause_ids) {
            return;
        }
        saw_history = true;
        let mut blocks = vec![text_block(
            render_borrowed_history_entry(entry, input.max_output_chars),
            false,
        )];
        append_borrowed_entry_image_blocks(entry, attachments, &mut blocks);
        messages.push(LlmMessage::new(
            borrowed_history_entry_llm_role(entry),
            blocks,
        ));
    });
    append_current_iteration_message(
        &mut messages,
        CurrentIterationMessageInput {
            saw_history,
            protocol_iteration: input.protocol_iteration,
            turn_causes: input.turn_causes,
            finalization: input.finalization,
            required_output: input.required_output,
            final_answer_format: input.final_answer_format,
            budget_suffix: input.budget_suffix,
        },
    );
    messages
}

#[cfg(test)]
pub(super) fn build_rlm_history_messages(
    input: RlmHistoryTestRenderInput<'_>,
    attachments: &mut Vec<LlmAttachment>,
) -> Vec<LlmMessage> {
    let mut messages = Vec::new();

    if !input.projection.entries().is_empty() {
        for entry in input.projection.entries() {
            let mut blocks = vec![text_block(
                render_history_entry(entry, input.max_output_chars),
                false,
            )];
            append_entry_image_blocks(entry, attachments, &mut blocks);
            messages.push(LlmMessage::new(history_entry_llm_role(entry), blocks));
        }
    }
    append_current_iteration_message(
        &mut messages,
        CurrentIterationMessageInput {
            saw_history: !input.projection.entries().is_empty(),
            protocol_iteration: input.protocol_iteration,
            turn_causes: &[],
            finalization: input.finalization,
            required_output: input.required_output,
            final_answer_format: input.final_answer_format,
            budget_suffix: input.budget_suffix,
        },
    );
    messages
}

fn append_current_iteration_message(
    messages: &mut Vec<LlmMessage>,
    input: CurrentIterationMessageInput<'_>,
) {
    if !input.saw_history {
        messages.push(LlmMessage::new(
            LlmRole::User,
            vec![text_block(
                "=== HISTORY ===\n\nNo chronological history is available.",
                false,
            )],
        ));
    } else {
        mark_last_history_text_cache_breakpoint(messages);
    }
    let mut current_prompt = format!(
        "\n\n\n=== CURRENT ITERATION: {} ===",
        input.protocol_iteration
    );
    if let Some(turn_events) = lash_core::render_turn_causes_prompt(input.turn_causes) {
        current_prompt.push_str("\n\n");
        current_prompt.push_str(&turn_events);
    }
    current_prompt.push_str("\n\n\n=== FINALIZATION ===\n\n");
    current_prompt.push_str(input.finalization);
    if let Some(block) = input.required_output {
        current_prompt.push_str("\n\n=== REQUIRED OUTPUT ===\n\n");
        current_prompt.push_str(block);
    }
    if let Some(guidance) = input.final_answer_format {
        current_prompt.push_str("\n\n=== FINAL ANSWER FORMAT ===\n\n");
        current_prompt.push_str(guidance);
    }
    if let Some(suffix) = input.budget_suffix {
        current_prompt.push_str("\n\n=== CONTEXT BUDGET ===\n\n");
        current_prompt.push_str(suffix);
    }
    messages.push(LlmMessage::new(
        LlmRole::User,
        vec![text_block(current_prompt, false)],
    ));
}

fn borrowed_entry_is_active_cause(
    entry: BorrowedChronologicalEntry<'_>,
    active_cause_ids: &HashSet<&str>,
) -> bool {
    matches!(
        entry.payload,
        BorrowedChronologicalPayload::Message(message)
            if matches!(message.role, lash_core::MessageRole::Event)
                && active_cause_ids.contains(message.id)
    )
}

fn text_block(text: impl Into<Arc<str>>, cache_breakpoint: bool) -> LlmContentBlock {
    LlmContentBlock::Text {
        text: text.into(),
        response_meta: None,
        cache_breakpoint,
    }
}

#[cfg(test)]
fn history_entry_llm_role(entry: &ChronologicalEntry) -> LlmRole {
    match &entry.payload {
        ChronologicalPayload::Message(message) => match message.role {
            lash_core::MessageRole::User => LlmRole::User,
            lash_core::MessageRole::Assistant => LlmRole::Assistant,
            lash_core::MessageRole::System => LlmRole::System,
            lash_core::MessageRole::Event => LlmRole::User,
        },
        ChronologicalPayload::ProtocolEvent(_) => LlmRole::Assistant,
    }
}

fn borrowed_history_entry_llm_role(entry: BorrowedChronologicalEntry<'_>) -> LlmRole {
    match entry.payload {
        BorrowedChronologicalPayload::Message(message) => match message.role {
            lash_core::MessageRole::User => LlmRole::User,
            lash_core::MessageRole::Assistant => LlmRole::Assistant,
            lash_core::MessageRole::System => LlmRole::System,
            lash_core::MessageRole::Event => LlmRole::User,
        },
        BorrowedChronologicalPayload::ProtocolEvent(_) => LlmRole::Assistant,
    }
}

fn mark_last_history_text_cache_breakpoint(messages: &mut [LlmMessage]) {
    for message in messages.iter_mut().rev() {
        let Some(blocks) = Arc::get_mut(&mut message.blocks) else {
            continue;
        };
        for block in blocks.iter_mut().rev() {
            if let LlmContentBlock::Text {
                text,
                cache_breakpoint,
                ..
            } = block
                && !text.trim().is_empty()
            {
                *cache_breakpoint = true;
                return;
            }
        }
    }
}

#[cfg(test)]
pub(super) fn render_history_prompt(history: &[RlmHistoryItem], max_output_chars: usize) -> String {
    if history.is_empty() {
        return "No chronological history is available.".to_string();
    }
    let mut rendered = String::new();
    for (index, item) in history.iter().enumerate() {
        if !rendered.is_empty() {
            rendered.push_str("\n\n");
        }
        rendered.push_str(&render_history_item(index, item, max_output_chars));
    }
    rendered
}

#[cfg(test)]
fn render_history_item(index: usize, item: &RlmHistoryItem, max_output_chars: usize) -> String {
    let mut rendered = String::new();
    match item {
        RlmHistoryItem::Message {
            id: _,
            role,
            content,
            attachments,
        } => append_history_message(
            &mut rendered,
            index,
            role,
            content,
            attachments,
            max_output_chars,
        ),
        RlmHistoryItem::LashlangStep {
            id: _,
            protocol_iteration,
            code,
            output,
            images,
            error,
            final_output,
        } => append_repl_step(
            &mut rendered,
            ReplStepRender {
                index,
                protocol_iteration: *protocol_iteration,
                code,
                output,
                images,
                error: error.as_deref(),
                final_output: final_output.as_ref(),
            },
        ),
    }
    rendered
}

#[cfg(test)]
fn render_history_entry(entry: &ChronologicalEntry, max_output_chars: usize) -> String {
    let mut rendered = String::new();
    match &entry.payload {
        ChronologicalPayload::Message(message) => {
            let content = message_history_text(message);
            let attachments = message
                .parts
                .iter()
                .filter_map(|part| {
                    let attachment = part.attachment.as_ref()?;
                    Some(RlmAttachmentRef {
                        id: part.id.clone(),
                        media_type: attachment.reference.media_type,
                        label: attachment.reference.label.clone(),
                        reference: attachment.reference.id.to_string(),
                    })
                })
                .collect::<Vec<_>>();
            append_history_message(
                &mut rendered,
                entry.index,
                &history_role(message.role),
                &content,
                &attachments,
                max_output_chars,
            );
        }
        ChronologicalPayload::ProtocolEvent(event) => {
            let Some(lash_rlm_types::RlmProtocolEvent::RlmTrajectoryEntry(step)) =
                decode_rlm_protocol_event(event)
            else {
                return rendered;
            };
            let images = step
                .images
                .iter()
                .map(|image| RlmImageRef {
                    id: image.id.to_string(),
                    media_type: image.media_type,
                    width: image.width,
                    height: image.height,
                    bytes: image.byte_len as usize,
                    label: image.label.clone(),
                })
                .collect::<Vec<_>>();
            append_repl_step(
                &mut rendered,
                ReplStepRender {
                    index: entry.index,
                    protocol_iteration: step.protocol_iteration,
                    code: &step.code,
                    output: &step.output,
                    images: &images,
                    error: step.error.as_deref(),
                    final_output: step.final_output.as_ref(),
                },
            );
        }
    }
    rendered
}

fn render_borrowed_history_entry(
    entry: BorrowedChronologicalEntry<'_>,
    max_output_chars: usize,
) -> String {
    let mut rendered = String::new();
    match entry.payload {
        BorrowedChronologicalPayload::Message(message) => {
            let content = message_history_text_parts(message.parts);
            let attachments = message
                .parts
                .iter()
                .filter_map(|part| {
                    let attachment = part.attachment.as_ref()?;
                    Some(RlmAttachmentRef {
                        id: part.id.clone(),
                        media_type: attachment.reference.media_type,
                        label: attachment.reference.label.clone(),
                        reference: attachment.reference.id.to_string(),
                    })
                })
                .collect::<Vec<_>>();
            append_history_message(
                &mut rendered,
                entry.index,
                &history_role(message.role),
                &content,
                &attachments,
                max_output_chars,
            );
        }
        BorrowedChronologicalPayload::ProtocolEvent(event) => {
            let Some(lash_rlm_types::RlmProtocolEvent::RlmTrajectoryEntry(step)) =
                decode_rlm_protocol_event(event)
            else {
                return rendered;
            };
            let images = step
                .images
                .iter()
                .map(|image| RlmImageRef {
                    id: image.id.to_string(),
                    media_type: image.media_type,
                    width: image.width,
                    height: image.height,
                    bytes: image.byte_len as usize,
                    label: image.label.clone(),
                })
                .collect::<Vec<_>>();
            append_repl_step(
                &mut rendered,
                ReplStepRender {
                    index: entry.index,
                    protocol_iteration: step.protocol_iteration,
                    code: &step.code,
                    output: &step.output,
                    images: &images,
                    error: step.error.as_deref(),
                    final_output: step.final_output.as_ref(),
                },
            );
        }
    }
    rendered
}

#[cfg(test)]
pub(super) fn append_entry_image_blocks(
    entry: &lash_core::ChronologicalEntry,
    attachments: &mut Vec<LlmAttachment>,
    blocks: &mut Vec<LlmContentBlock>,
) {
    match &entry.payload {
        lash_core::ChronologicalPayload::Message(message) => {
            for part in message.parts.iter() {
                let Some(attachment) = part.attachment.as_ref() else {
                    continue;
                };
                let attachment_idx = attachments.len();
                attachments.push(LlmAttachment::reference(attachment.reference.clone()));
                blocks.push(LlmContentBlock::Image { attachment_idx });
            }
        }
        lash_core::ChronologicalPayload::ProtocolEvent(event) => {
            if let Some(lash_rlm_types::RlmProtocolEvent::RlmTrajectoryEntry(entry)) =
                decode_rlm_protocol_event(event)
            {
                for image in &entry.images {
                    let attachment_idx = attachments.len();
                    attachments.push(LlmAttachment::reference(image.clone()));
                    blocks.push(LlmContentBlock::Image { attachment_idx });
                }
            }
        }
    }
}

fn append_borrowed_entry_image_blocks(
    entry: BorrowedChronologicalEntry<'_>,
    attachments: &mut Vec<LlmAttachment>,
    blocks: &mut Vec<LlmContentBlock>,
) {
    match entry.payload {
        BorrowedChronologicalPayload::Message(message) => {
            for part in message.parts {
                let Some(attachment) = part.attachment.as_ref() else {
                    continue;
                };
                let attachment_idx = attachments.len();
                attachments.push(LlmAttachment::reference(attachment.reference.clone()));
                blocks.push(LlmContentBlock::Image { attachment_idx });
            }
        }
        BorrowedChronologicalPayload::ProtocolEvent(event) => {
            if let Some(lash_rlm_types::RlmProtocolEvent::RlmTrajectoryEntry(entry)) =
                decode_rlm_protocol_event(event)
            {
                for image in &entry.images {
                    let attachment_idx = attachments.len();
                    attachments.push(LlmAttachment::reference(image.clone()));
                    blocks.push(LlmContentBlock::Image { attachment_idx });
                }
            }
        }
    }
}

fn append_history_message(
    out: &mut String,
    index: usize,
    role: &RlmHistoryRole,
    content: &str,
    attachments: &[RlmAttachmentRef],
    max_output_chars: usize,
) {
    let (preview, raw_len) = head_tail_truncate(content, max_output_chars);
    let full_ref = truncated_ref(
        raw_len,
        max_output_chars,
        &format!("history[{index}].content"),
    );
    let _ = write!(
        out,
        "--- history[{index}] · {} message · {raw_len} chars{full_ref} ---\n\n{preview}",
        message_role_label(role).to_lowercase(),
    );
    if !attachments.is_empty() {
        out.push_str("\n\nAttachments:");
        for (attachment_index, attachment) in attachments.iter().enumerate() {
            let rendered = serde_json::to_string(attachment)
                .unwrap_or_else(|_| "{\"error\":\"unrenderable attachment\"}".to_string());
            let _ = write!(
                out,
                "\n- history[{index}].attachments[{attachment_index}]: {rendered}"
            );
        }
    }
}

#[cfg(test)]
fn message_history_text(message: &lash_core::Message) -> String {
    message_history_text_parts(message.parts.as_slice())
}

fn message_history_text_parts(parts: &[lash_core::Part]) -> String {
    let chunks = parts
        .iter()
        .filter(|part| {
            matches!(
                part.kind,
                lash_core::PartKind::Text | lash_core::PartKind::Prose
            )
        })
        .map(|part| part.content.trim())
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>();
    chunks.join("\n\n")
}

fn history_role(role: lash_core::MessageRole) -> RlmHistoryRole {
    match role {
        lash_core::MessageRole::User => RlmHistoryRole::User,
        lash_core::MessageRole::System => RlmHistoryRole::System,
        lash_core::MessageRole::Assistant => RlmHistoryRole::Assistant,
        lash_core::MessageRole::Event => RlmHistoryRole::Event,
    }
}

fn message_role_label(role: &RlmHistoryRole) -> &'static str {
    match role {
        RlmHistoryRole::User => "User",
        RlmHistoryRole::Assistant => "Assistant",
        RlmHistoryRole::System => "System",
        RlmHistoryRole::Event => "Event",
    }
}

struct ReplStepRender<'a> {
    index: usize,
    protocol_iteration: usize,
    code: &'a str,
    output: &'a [String],
    images: &'a [RlmImageRef],
    error: Option<&'a str>,
    final_output: Option<&'a serde_json::Value>,
}

fn append_repl_step(out: &mut String, step: ReplStepRender<'_>) {
    let ReplStepRender {
        index,
        protocol_iteration,
        code,
        output,
        images,
        error,
        final_output,
    } = step;
    let _ = write!(
        out,
        "--- history[{index}] · lashlang step · protocol_iteration {protocol_iteration} ---\n\nCode:\n{}",
        indent_source(code.trim()),
    );

    // One block per `print` (or raw stdout emission). Tool calls used to
    // get their own section here for retrieval-without-print, but that
    // duplicated content the model could fetch via `print result` and
    // bloated history; the `code` field above shows every receiver
    // operation the model wrote.
    for (output_index, item) in output.iter().enumerate() {
        let (preview, projected_lossy) = project_history_output(item);
        let raw_len = item.chars().count();
        let full_ref = projected_ref(
            projected_lossy,
            &format!("history[{index}].output[{output_index}]"),
        );
        let _ = write!(
            out,
            "\n\nhistory[{index}].output[{output_index}] ({raw_len} chars{full_ref}):\n{preview}"
        );
    }

    if !images.is_empty() {
        out.push_str("\n\nImages:");
        for (image_index, image) in images.iter().enumerate() {
            let rendered = serde_json::to_string(image)
                .unwrap_or_else(|_| "{\"error\":\"unrenderable image\"}".to_string());
            let _ = write!(
                out,
                "\n- history[{index}].images[{image_index}]: {rendered}"
            );
        }
    }

    if let Some(error) = error {
        out.push_str("\n\nError:\n");
        out.push_str(error);
    }
    if let Some(final_output) = final_output {
        out.push_str("\n\nFinal output:\n");
        out.push_str(
            &serde_json::to_string_pretty(final_output)
                .unwrap_or_else(|_| final_output.to_string()),
        );
    }
}

fn truncated_ref(raw_len: usize, max_output_chars: usize, reference: &str) -> String {
    if raw_len > max_output_chars {
        format!(", full: {reference}")
    } else {
        String::new()
    }
}

fn projected_ref(projected_lossy: bool, reference: &str) -> String {
    if projected_lossy {
        format!(", full: {reference}")
    } else {
        String::new()
    }
}

fn indent_source(source: &str) -> String {
    if source.is_empty() {
        return "    (empty)".to_string();
    }
    source
        .lines()
        .map(|line| format!("    {line}"))
        .collect::<Vec<_>>()
        .join("\n")
}

fn project_history_output(item: &str) -> (String, bool) {
    let value = history_output_value(item);
    let projected = crate::rlm_support::print_history_projector()
        .project_blocking(ValueProjectionContext::new(&value));
    let lossy = projection_is_lossy(item, &projected);
    (projected, lossy)
}

fn history_output_value(item: &str) -> FlowValue {
    let trimmed = item.trim_start();
    if (trimmed.starts_with('{') || trimmed.starts_with('['))
        && let Ok(value) = serde_json::from_str::<serde_json::Value>(item)
    {
        return json_to_flow_value(value);
    }
    FlowValue::String(item.into())
}

fn projection_is_lossy(original: &str, projected: &str) -> bool {
    projected.contains("truncated")
        || projected.contains("omitted")
        || projected.contains("max depth")
        || projected.chars().count() < original.chars().count()
}