plana-text 0.4.2

LLM text accumulation and structured streaming abstractions
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use super::{InputMessage, StreamChunkKind, StreamSegment};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmStream {
    #[serde(default)]
    pub(crate) segments: Vec<StreamSegment>,
}

/// Trim a notice source; `None` when it carries no name at all (the row
/// exists only to name an origin, so an anonymous one is refused).
fn normalise_source(source: &str) -> Option<String> {
    let trimmed = source.trim();
    (!trimmed.is_empty()).then(|| trimmed.to_string())
}

/// Trim a notice detail; `None` when blank (an empty second line is
/// noise, not information).
fn normalise_detail(detail: Option<&str>) -> Option<String> {
    detail
        .map(str::trim)
        .filter(|d| !d.is_empty())
        .map(str::to_string)
}

impl LlmStream {
    pub fn empty() -> Self {
        Self {
            segments: Vec::new(),
        }
    }

    /// Prepend the composed INPUT side as the stream's leading segment
    /// (user direction 2026-10-10): the caller — scepter's pipeline,
    /// where the composed messages live — records what the model was
    /// actually fed (the system stack + the task turn) so the run
    /// record's input lane carries the prompt composition. One call per
    /// captured stream: the first round's composition; later rounds
    /// resend the same stack plus segments the stream already records.
    pub fn prepend_input(&mut self, messages: Vec<InputMessage>) {
        if messages.is_empty() {
            return;
        }
        self.segments.insert(
            0,
            StreamSegment::Input {
                messages,
                message_id: None,
                ts: None,
            },
        );
    }

    /// Append a PROMPT NOTICE at the stream's tail — the position for a
    /// change observed AFTER the output it did not shape (a retry's
    /// recomposition, an operator edit mid-settle).
    ///
    /// Use [`Self::push_notice_after_input`] instead when the notice
    /// describes THIS stream's own prompt: there the change took effect
    /// BEFORE the stream's output, and the tail slot would read as if the
    /// output preceded the change.
    ///
    /// Returns whether the notice was recorded: a blank `source` is a bug
    /// at the call site (the row exists only to name the origin), so it is
    /// dropped rather than recorded as an anonymous notice.
    pub fn push_notice(&mut self, source: &str, detail: Option<&str>, ts: Option<u64>) -> bool {
        let Some(source) = normalise_source(source) else {
            return false;
        };
        self.segments.push(StreamSegment::PromptNotice {
            source,
            detail: normalise_detail(detail),
            ts,
        });
        true
    }

    /// Append a PROMPT NOTICE directly AFTER the leading `Input` segment —
    /// the position for a notice that describes THIS stream's prompt: the
    /// rows under it are the output that prompt produced, so the
    /// provenance line leads them instead of trailing them. Without a
    /// leading `Input` segment the notice leads the stream (index 0).
    ///
    /// Same blank-source contract as [`Self::push_notice`].
    pub fn push_notice_after_input(
        &mut self,
        source: &str,
        detail: Option<&str>,
        ts: Option<u64>,
    ) -> bool {
        let Some(source) = normalise_source(source) else {
            return false;
        };
        let pos = usize::from(
            self.segments
                .first()
                .is_some_and(|seg| matches!(seg, StreamSegment::Input { .. })),
        );
        self.segments.insert(
            pos,
            StreamSegment::PromptNotice {
                source,
                detail: normalise_detail(detail),
                ts,
            },
        );
        true
    }

    pub fn from_plain_text(text: &str) -> Self {
        let segments = if text.is_empty() {
            Vec::new()
        } else {
            vec![StreamSegment::Text {
                text: text.to_string(),
                message_id: None,
                ts: None,
            }]
        };
        Self { segments }
    }

    pub fn raw_text(&self) -> String {
        self.segments.iter().map(|s| s.text_or_json()).collect()
    }

    pub fn segments(&self) -> &[StreamSegment] {
        &self.segments
    }

    pub fn segments_mut(&mut self) -> &mut [StreamSegment] {
        &mut self.segments
    }

    pub fn patch_tool_result_for_call(
        &mut self,
        call_id: Uuid,
        result_data: serde_json::Value,
        success: bool,
        duration_ms: Option<u64>,
    ) -> bool {
        let sealed_marker = "[stream sealed before result arrived]";
        let sealed_value = serde_json::Value::String(sealed_marker.to_string());

        for seg in &mut self.segments {
            if let StreamSegment::ToolResult {
                call_id: seg_call_id,
                data,
                success: seg_success,
                duration_ms: seg_duration,
                ..
            } = seg
                && *seg_call_id == call_id
                && *data == sealed_value
            {
                *data = result_data;
                *seg_success = success;
                *seg_duration = duration_ms;
                return true;
            }
        }

        let has_call = self
            .segments
            .iter()
            .any(|s| matches!(s, StreamSegment::ToolCall { call_id: cid, .. } if *cid == call_id));
        if has_call {
            let tool_name = self
                .segments
                .iter()
                .find_map(|s| match s {
                    StreamSegment::ToolCall {
                        call_id: cid,
                        tool_name,
                        ..
                    } if *cid == call_id => Some(tool_name.clone()),
                    _ => None,
                })
                .unwrap_or_default();
            let agent_type = self
                .segments
                .iter()
                .find_map(|s| match s {
                    StreamSegment::ToolCall {
                        call_id: cid,
                        agent_type,
                        ..
                    } if *cid == call_id => Some(agent_type.clone()),
                    _ => None,
                })
                .flatten();

            let pos = self
                .segments
                .iter()
                .rposition(
                    |s| matches!(s, StreamSegment::ToolCall { call_id: cid, .. } if *cid == call_id),
                )
                .map(|i| i + 1)
                .unwrap_or(self.segments.len());
            self.segments.insert(
                pos,
                StreamSegment::ToolResult {
                    tool_name,
                    call_id,
                    success,
                    data: result_data,
                    duration_ms,
                    agent_type,
                    message_id: None,
                    ts: None,
                },
            );
            return true;
        }

        false
    }

    pub fn has_tool_call_with_id(&self, call_id: Uuid) -> bool {
        self.segments
            .iter()
            .any(|s| matches!(s, StreamSegment::ToolCall { call_id: cid, .. } if *cid == call_id))
    }

    pub fn is_empty(&self) -> bool {
        self.segments.is_empty()
            || self.segments.iter().all(|s| match s {
                StreamSegment::Text { text, .. }
                | StreamSegment::Thinking { text, .. }
                | StreamSegment::DeepThinking { text, .. } => text.is_empty(),
                StreamSegment::ToolCall { params, .. } => params.is_null(),
                StreamSegment::ToolResult { data, .. } => data.is_null(),
                // An Input segment always carries its messages' payload —
                // a stream holding one is not empty.
                StreamSegment::Input { .. } => false,
                // A prompt notice is a recorded event of its own: a
                // stream holding only the notice is still an event the
                // panel must render, never "no output".
                StreamSegment::PromptNotice { .. } => false,
            })
    }

    pub fn text_segments(&self) -> impl Iterator<Item = (StreamChunkKind, &str)> {
        self.segments.iter().filter_map(|seg| {
            let kind = seg.chunk_kind()?;
            Some((kind, seg.text()))
        })
    }

    pub fn tool_segments(&self) -> impl Iterator<Item = &StreamSegment> {
        self.segments.iter().filter(|s| s.is_tool())
    }

    pub fn thinking_text(&self) -> String {
        self.text_segments()
            .filter(|(kind, _)| {
                matches!(
                    kind,
                    StreamChunkKind::Thinking | StreamChunkKind::DeepThinking
                )
            })
            .map(|(_, text)| text)
            .collect::<Vec<_>>()
            .join("")
    }

    pub fn display_text(&self) -> String {
        self.text_segments()
            .filter(|(kind, _)| matches!(kind, StreamChunkKind::Text))
            .map(|(_, text)| text)
            .collect::<Vec<_>>()
            .join("")
    }

    pub fn len(&self) -> usize {
        self.segments.iter().map(|s| s.text_or_json().len()).sum()
    }
}

impl std::fmt::Display for LlmStream {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for seg in &self.segments {
            f.write_str(&seg.text_or_json())?;
        }
        Ok(())
    }
}

impl Default for LlmStream {
    fn default() -> Self {
        Self::empty()
    }
}

impl From<String> for LlmStream {
    fn from(s: String) -> Self {
        Self::from_plain_text(&s)
    }
}

impl From<&str> for LlmStream {
    fn from(s: &str) -> Self {
        Self::from_plain_text(s)
    }
}

impl PartialEq for LlmStream {
    fn eq(&self, other: &Self) -> bool {
        self.segments == other.segments
    }
}

impl Eq for LlmStream {}