lash_sansio/session_model/
message.rs

1use crate::AttachmentRef;
2use crate::llm::types::{
3    LlmAttachment, LlmContentBlock, LlmMessage, LlmRole, ProviderReasoningReplay,
4    ProviderReplayMeta, ResponseTextMeta,
5};
6use std::collections::HashSet;
7use std::sync::{Arc, OnceLock};
8
9// ─── Structured message types for context-aware pruning ───
10
11/// A structured message with typed parts for context management.
12///
13/// `parts` is `Arc`-shared so cloning a `Message` is one Arc bump per
14/// message field rather than a deep-clone of every `Part`. Construct with
15/// `parts: shared_parts(vec![...])` or `parts: Arc::new(...)`. Mutate via
16/// `Arc::make_mut(&mut message.parts)` when truly needed; most plugin
17/// pipelines should produce a fresh `Vec<Part>` and assign it.
18#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
19pub struct Message {
20    pub id: String,
21    pub role: MessageRole,
22    pub parts: Arc<Vec<Part>>,
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub origin: Option<MessageOrigin>,
25}
26
27/// Wrap a `Vec<Part>` for the `Message::parts` field. Use this in struct
28/// literals and tests (`parts: shared_parts(vec![Part { ... }])`) so the
29/// call sites stay short and uniform.
30#[inline]
31pub fn shared_parts(parts: Vec<Part>) -> Arc<Vec<Part>> {
32    Arc::new(parts)
33}
34
35#[derive(Clone, Copy, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
36pub enum MessageRole {
37    User,
38    Assistant,
39    System,
40}
41
42#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
43#[serde(tag = "kind", rename_all = "snake_case")]
44pub enum MessageOrigin {
45    Plugin {
46        plugin_id: String,
47        #[serde(default, skip_serializing_if = "std::ops::Not::not")]
48        transient: bool,
49    },
50}
51
52#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
53pub struct Part {
54    /// e.g. "m3.p0"
55    pub id: String,
56    pub kind: PartKind,
57    pub content: String,
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub attachment: Option<PartAttachment>,
60    #[serde(default, skip_serializing_if = "Option::is_none")]
61    pub tool_call_id: Option<String>,
62    #[serde(default, skip_serializing_if = "Option::is_none")]
63    pub tool_name: Option<String>,
64    /// Opaque provider replay state attached to a `ToolCall` part.
65    #[serde(default, skip_serializing_if = "Option::is_none")]
66    pub tool_replay: Option<ProviderReplayMeta>,
67    pub prune_state: PruneState,
68    /// Populated only for `PartKind::Reasoning` parts. Carries opaque
69    /// provider replay metadata so the adapter can re-emit the exact same
70    /// reasoning item on subsequent turns.
71    /// `#[serde(default, skip_serializing_if)]` so older snapshots that
72    /// predate this field round-trip unchanged.
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub reasoning_meta: Option<ProviderReasoningReplay>,
75    /// Provider message metadata for assistant text parts. Legacy snapshots
76    /// omit it; adapters synthesize deterministic ids when replaying older
77    /// assistant text.
78    #[serde(default, skip_serializing_if = "Option::is_none")]
79    pub response_meta: Option<ResponseTextMeta>,
80}
81
82#[derive(Clone, Copy, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
83pub enum PartKind {
84    Text,
85    Image,
86    Code,
87    Output,
88    Error,
89    Prose,
90    ToolCall,
91    ToolResult,
92    /// Chain-of-thought / reasoning item captured from providers that expose
93    /// a reasoning channel. `content` holds the human-readable summary for
94    /// display (fix 1.3a). The encrypted blob and raw `summary`/`id` needed
95    /// to re-feed the model on the next turn (fix 1.3b) live in
96    /// `reasoning_meta`. Reasoning parts are preserved across snapshots so
97    /// next-turn re-feeding survives session resume; they are never rendered
98    /// into the flat chat prompt. Provider adapters decide whether and how
99    /// to re-emit them through their native channel.
100    Reasoning,
101}
102
103#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
104pub struct PartAttachment {
105    pub reference: AttachmentRef,
106}
107
108#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
109pub enum PruneState {
110    Intact,
111    Cleared,
112    Deleted {
113        breadcrumb: String,
114        archive_hash: String,
115    },
116    Summarized {
117        summary: String,
118        archive_hash: String,
119    },
120}
121
122impl Part {
123    pub fn prompt_char_count(&self) -> usize {
124        // Reasoning parts are not user-visible text and aren't sent to the
125        // model as flat prompt content. Provider adapters may re-emit them
126        // via structured replay metadata instead. Excluding them from the
127        // accounting keeps the rolling-history plugin's prune decisions
128        // driven by real conversation content.
129        if matches!(self.kind, PartKind::Reasoning) {
130            return 0;
131        }
132        if matches!(self.kind, PartKind::Image) {
133            return self
134                .attachment
135                .as_ref()
136                .map(|attachment| attachment.reference.id.as_str().len())
137                .unwrap_or_else(|| self.render().len());
138        }
139        self.render().len()
140    }
141
142    pub(crate) fn render(&self) -> String {
143        if matches!(self.kind, PartKind::Image) {
144            return if self.attachment.is_some() || self.content.trim().is_empty() {
145                "[Image attached]".to_string()
146            } else {
147                self.content.clone()
148            };
149        }
150        match &self.prune_state {
151            PruneState::Intact => self.content.clone(),
152            PruneState::Cleared => "[Old tool result content cleared]".to_string(),
153            PruneState::Deleted {
154                breadcrumb,
155                archive_hash,
156            } => format!("[pruned:{} — {}]", archive_hash, breadcrumb),
157            PruneState::Summarized {
158                summary,
159                archive_hash,
160            } => format!("[SUMMARY of original {}]\n{}", archive_hash, summary),
161        }
162    }
163}
164
165impl Message {
166    /// Total character count of all parts (rendered).
167    pub fn char_count(&self) -> usize {
168        self.parts.iter().map(Part::prompt_char_count).sum()
169    }
170
171    pub fn is_transient(&self) -> bool {
172        matches!(
173            self.origin,
174            Some(MessageOrigin::Plugin {
175                transient: true,
176                ..
177            })
178        )
179    }
180}
181
182fn render_part_for_chat(role: MessageRole, part: &Part) -> String {
183    let rendered = part.render();
184    match role {
185        MessageRole::System => match part.kind {
186            PartKind::Code => rendered,
187            PartKind::Output => format!("<output>\n{}\n</output>", rendered),
188            PartKind::Error => format!("<error>\n{}\n</error>", rendered),
189            PartKind::Text
190            | PartKind::Image
191            | PartKind::Prose
192            | PartKind::ToolCall
193            | PartKind::ToolResult
194            | PartKind::Reasoning => rendered,
195        },
196        MessageRole::Assistant => match part.kind {
197            PartKind::Code => rendered,
198            PartKind::ToolCall => render_assistant_tool_call(part, &rendered),
199            PartKind::Prose | PartKind::Text | PartKind::Image | PartKind::ToolResult => rendered,
200            PartKind::Reasoning => rendered,
201            _ => rendered,
202        },
203        MessageRole::User => rendered,
204    }
205}
206
207fn render_assistant_tool_call(part: &Part, rendered: &str) -> String {
208    let tool_name = part.tool_name.as_deref().unwrap_or("tool");
209    let trimmed = rendered.trim();
210    if trimmed.is_empty() || trimmed == "{}" {
211        format!("{tool_name}()")
212    } else {
213        format!("{tool_name}({trimmed})")
214    }
215}
216
217fn attachment_from_part(part: &Part) -> Option<LlmAttachment> {
218    if !matches!(part.kind, PartKind::Image) {
219        return None;
220    }
221    let attachment = part.attachment.as_ref()?;
222    Some(LlmAttachment::reference(attachment.reference.clone()))
223}
224
225fn render_message_for_transcript(msg: &Message, attachments: &mut Vec<LlmAttachment>) -> String {
226    let mut out = Vec::new();
227    for part in msg.parts.iter() {
228        // Reasoning items are display-only from the transcript's point of
229        // view — they are never replayed as flat text. Provider adapters use
230        // structured replay metadata when they can re-emit reasoning.
231        if matches!(part.kind, PartKind::Reasoning) {
232            continue;
233        }
234        if let Some(attachment) = attachment_from_part(part) {
235            attachments.push(attachment);
236            out.push("[Image attached]".to_string());
237            continue;
238        }
239        let rendered = render_part_for_chat(msg.role, part);
240        if !rendered.trim().is_empty() {
241            out.push(rendered);
242        }
243    }
244    out.join("\n\n")
245}
246
247#[derive(Clone, Debug, Default, PartialEq, Eq)]
248pub struct RenderedPrompt {
249    pub messages: Vec<LlmMessage>,
250    pub attachments: Vec<LlmAttachment>,
251}
252
253/// Memoized render of a `MessageSequence`'s `base`. Shared across the
254/// per-iteration `MessageSequence` instances that wrap the same base
255/// (typically the `SessionGraphCache`'s projected messages) so the
256/// chat projector's `render_prompt` walk happens once per turn instead
257/// of once per LLM iteration.
258pub type BaseRenderCache = OnceLock<RenderedPrompt>;
259
260#[derive(Debug)]
261pub struct MessageSequence {
262    base: Arc<Vec<Message>>,
263    delta: Vec<Message>,
264    owned: Option<Vec<Message>>,
265    materialized: OnceLock<Arc<Vec<Message>>>,
266    base_rendered: Option<Arc<BaseRenderCache>>,
267}
268
269impl Clone for MessageSequence {
270    fn clone(&self) -> Self {
271        Self {
272            base: Arc::clone(&self.base),
273            delta: self.delta.clone(),
274            owned: self.owned.clone(),
275            materialized: OnceLock::new(),
276            base_rendered: self.base_rendered.as_ref().map(Arc::clone),
277        }
278    }
279}
280
281impl Default for MessageSequence {
282    fn default() -> Self {
283        Self::from_owned(Vec::new())
284    }
285}
286
287impl From<Vec<Message>> for MessageSequence {
288    fn from(messages: Vec<Message>) -> Self {
289        Self::from_owned(messages)
290    }
291}
292
293impl std::ops::Deref for MessageSequence {
294    type Target = [Message];
295
296    fn deref(&self) -> &Self::Target {
297        self.as_slice()
298    }
299}
300
301impl MessageSequence {
302    pub fn from_owned(messages: Vec<Message>) -> Self {
303        Self {
304            base: Arc::new(Vec::new()),
305            delta: Vec::new(),
306            owned: Some(messages),
307            materialized: OnceLock::new(),
308            base_rendered: None,
309        }
310    }
311
312    pub fn from_base(base: Arc<Vec<Message>>) -> Self {
313        Self {
314            base,
315            delta: Vec::new(),
316            owned: None,
317            materialized: OnceLock::new(),
318            base_rendered: None,
319        }
320    }
321
322    pub fn from_base_and_delta(base: Arc<Vec<Message>>, delta: Vec<Message>) -> Self {
323        Self {
324            base,
325            delta,
326            owned: None,
327            materialized: OnceLock::new(),
328            base_rendered: None,
329        }
330    }
331
332    /// Attach a shared render cache for the `base` portion. Subsequent
333    /// `render_prompt` calls will reuse the memoized `RenderedPrompt` for
334    /// the base instead of rewalking it. The delta is always re-rendered
335    /// because it changes per LLM iteration. Returns `self` for chaining.
336    pub fn with_base_render_cache(mut self, cache: Arc<BaseRenderCache>) -> Self {
337        self.base_rendered = Some(cache);
338        self
339    }
340
341    pub fn len(&self) -> usize {
342        match &self.owned {
343            Some(owned) => owned.len(),
344            None => self.base.len() + self.delta.len(),
345        }
346    }
347
348    pub fn is_empty(&self) -> bool {
349        self.len() == 0
350    }
351
352    pub fn iter(&self) -> MessageSequenceIter<'_> {
353        match self.owned.as_ref() {
354            Some(owned) => MessageSequenceIter::Owned(owned.iter()),
355            None => MessageSequenceIter::Split(self.base.iter().chain(self.delta.iter())),
356        }
357    }
358
359    pub fn as_slice(&self) -> &[Message] {
360        if let Some(owned) = &self.owned {
361            return owned.as_slice();
362        }
363        if self.delta.is_empty() {
364            return self.base.as_slice();
365        }
366        self.materialized
367            .get_or_init(|| {
368                let mut combined = Vec::with_capacity(self.base.len() + self.delta.len());
369                combined.extend(self.base.iter().cloned());
370                combined.extend(self.delta.iter().cloned());
371                Arc::new(combined)
372            })
373            .as_slice()
374    }
375
376    pub fn shared(&self) -> Arc<Vec<Message>> {
377        if let Some(owned) = &self.owned {
378            return Arc::clone(self.materialized.get_or_init(|| Arc::new(owned.clone())));
379        }
380        if self.delta.is_empty() {
381            return Arc::clone(&self.base);
382        }
383        Arc::clone(self.materialized.get_or_init(|| {
384            let mut combined = Vec::with_capacity(self.base.len() + self.delta.len());
385            combined.extend(self.base.iter().cloned());
386            combined.extend(self.delta.iter().cloned());
387            Arc::new(combined)
388        }))
389    }
390
391    pub fn make_mut(&mut self) -> &mut Vec<Message> {
392        if self.owned.is_none() {
393            let owned = if self.delta.is_empty() {
394                Arc::unwrap_or_clone(Arc::clone(&self.base))
395            } else if let Some(materialized) = self.materialized.get() {
396                Arc::unwrap_or_clone(Arc::clone(materialized))
397            } else {
398                let mut combined = Vec::with_capacity(self.base.len() + self.delta.len());
399                combined.extend(self.base.iter().cloned());
400                combined.extend(self.delta.iter().cloned());
401                combined
402            };
403            self.owned = Some(owned);
404            self.base = Arc::new(Vec::new());
405            self.delta.clear();
406        }
407        self.materialized = OnceLock::new();
408        self.owned.as_mut().expect("message sequence owned state")
409    }
410
411    pub fn push(&mut self, message: Message) {
412        if let Some(owned) = self.owned.as_mut() {
413            owned.push(message);
414        } else {
415            self.delta.push(message);
416        }
417        self.materialized = OnceLock::new();
418    }
419
420    pub fn extend(&mut self, messages: Vec<Message>) {
421        if messages.is_empty() {
422            return;
423        }
424        if let Some(owned) = self.owned.as_mut() {
425            owned.extend(messages);
426        } else {
427            self.delta.extend(messages);
428        }
429        self.materialized = OnceLock::new();
430    }
431
432    pub fn replace(&mut self, messages: Vec<Message>) {
433        self.base = Arc::new(Vec::new());
434        self.delta.clear();
435        self.owned = Some(messages);
436        self.materialized = OnceLock::new();
437    }
438
439    pub fn into_vec(self) -> Vec<Message> {
440        if let Some(owned) = self.owned {
441            return owned;
442        }
443        if self.delta.is_empty() {
444            return Arc::unwrap_or_clone(self.base);
445        }
446        if let Some(materialized) = self.materialized.into_inner() {
447            return Arc::unwrap_or_clone(materialized);
448        }
449        let mut combined = Vec::with_capacity(self.base.len() + self.delta.len());
450        combined.extend(self.base.iter().cloned());
451        combined.extend(self.delta);
452        combined
453    }
454
455    pub fn render_prompt(&self) -> RenderedPrompt {
456        if let Some(owned) = &self.owned {
457            return render_prompt(owned.as_slice());
458        }
459        if self.base.is_empty() {
460            return render_prompt(self.delta.as_slice());
461        }
462        let mut rendered = match &self.base_rendered {
463            Some(cache) => cache
464                .get_or_init(|| render_prompt(self.base.as_slice()))
465                .clone(),
466            None => render_prompt(self.base.as_slice()),
467        };
468        if !self.delta.is_empty() {
469            append_rendered_prompt(&mut rendered, self.delta.as_slice());
470        }
471        rendered
472    }
473}
474
475pub enum MessageSequenceIter<'a> {
476    Owned(std::slice::Iter<'a, Message>),
477    Split(std::iter::Chain<std::slice::Iter<'a, Message>, std::slice::Iter<'a, Message>>),
478}
479
480impl<'a> Iterator for MessageSequenceIter<'a> {
481    type Item = &'a Message;
482
483    fn next(&mut self) -> Option<Self::Item> {
484        match self {
485            Self::Owned(iter) => iter.next(),
486            Self::Split(iter) => iter.next(),
487        }
488    }
489}
490
491#[derive(Clone, Debug, Default)]
492struct TranscriptTurn {
493    user: Vec<String>,
494    assistant: Vec<String>,
495}
496
497pub fn render_prompt(msgs: &[Message]) -> RenderedPrompt {
498    let mut rendered = RenderedPrompt::default();
499    append_rendered_prompt(&mut rendered, msgs);
500    rendered
501}
502
503pub fn messages_are_prompt_resume_safe<'a>(
504    messages: impl IntoIterator<Item = &'a Message>,
505) -> bool {
506    let mut seen_tool_calls = HashSet::new();
507    let mut completed_tool_calls = HashSet::new();
508
509    for message in messages {
510        for part in message.parts.iter() {
511            // Reasoning parts don't participate in tool pairing and are
512            // always safe to resume through.
513            if matches!(part.kind, PartKind::Reasoning) {
514                continue;
515            }
516            match part.kind {
517                PartKind::ToolCall => {
518                    if !matches!(message.role, MessageRole::Assistant) {
519                        return false;
520                    }
521                    let Some(call_id) = part
522                        .tool_call_id
523                        .as_deref()
524                        .map(str::trim)
525                        .filter(|call_id| !call_id.is_empty())
526                    else {
527                        return false;
528                    };
529                    if !seen_tool_calls.insert(call_id) {
530                        return false;
531                    }
532                }
533                PartKind::ToolResult => {
534                    if !matches!(message.role, MessageRole::User) {
535                        return false;
536                    }
537                    let Some(call_id) = part
538                        .tool_call_id
539                        .as_deref()
540                        .map(str::trim)
541                        .filter(|call_id| !call_id.is_empty())
542                    else {
543                        return false;
544                    };
545                    if !seen_tool_calls.contains(call_id) {
546                        return false;
547                    }
548                    if !completed_tool_calls.insert(call_id) {
549                        return false;
550                    }
551                }
552                _ => {}
553            }
554        }
555    }
556
557    seen_tool_calls.len() == completed_tool_calls.len()
558}
559
560pub fn render_transcript_prompt(msgs: &[Message]) -> RenderedPrompt {
561    let mut attachments = Vec::new();
562    let mut turns = Vec::new();
563    let mut current = TranscriptTurn::default();
564    let mut has_current = false;
565
566    for msg in msgs {
567        let text = render_message_for_transcript(msg, &mut attachments);
568        let has_text = !text.trim().is_empty();
569        match msg.role {
570            MessageRole::User => {
571                if has_current && (!current.user.is_empty() || !current.assistant.is_empty()) {
572                    turns.push(current);
573                    current = TranscriptTurn::default();
574                }
575                if has_text {
576                    current.user.push(text);
577                }
578                has_current = true;
579            }
580            MessageRole::Assistant | MessageRole::System => {
581                if !has_current {
582                    has_current = true;
583                }
584                if has_text {
585                    current.assistant.push(text);
586                }
587            }
588        }
589    }
590
591    if has_current && (!current.user.is_empty() || !current.assistant.is_empty()) {
592        turns.push(current);
593    }
594
595    let mut text = String::new();
596    text.push_str(
597        "History:\nThis is a chronological transcript. `Assistant` refers to Lash, and you are continuing the same session.\n\n",
598    );
599    for (idx, turn) in turns.iter().enumerate() {
600        text.push_str(&format!("=== Turn {} ===\n", idx + 1));
601        text.push_str("User:\n");
602        if turn.user.is_empty() {
603            text.push_str("[No user content recorded]\n");
604        } else {
605            text.push_str(&turn.user.join("\n\n"));
606            text.push('\n');
607        }
608        text.push('\n');
609        text.push_str("Assistant (Lash, continuing this transcript):\n");
610        let is_current_pending_turn = idx + 1 == turns.len() && turn.assistant.is_empty();
611        if turn.assistant.is_empty() && !is_current_pending_turn {
612            text.push_str("[No assistant content recorded]\n");
613        } else if !turn.assistant.is_empty() {
614            text.push_str(&turn.assistant.join("\n\n"));
615            text.push('\n');
616        }
617        text.push('\n');
618    }
619    text.push_str(
620        "Continue from the latest turn as Lash.\nIf the task is complete, provide the final answer.\nOtherwise produce the next valid step for this runtime.",
621    );
622
623    RenderedPrompt {
624        messages: vec![LlmMessage::text(LlmRole::User, text)],
625        attachments,
626    }
627}
628
629pub fn append_rendered_prompt(rendered: &mut RenderedPrompt, msgs: &[Message]) {
630    append_structured_prompt(rendered, msgs)
631}
632
633#[cfg(test)]
634fn render_structured_prompt(msgs: &[Message]) -> RenderedPrompt {
635    let mut rendered = RenderedPrompt::default();
636    append_structured_prompt(&mut rendered, msgs);
637    rendered
638}
639
640fn append_structured_prompt(rendered: &mut RenderedPrompt, msgs: &[Message]) {
641    for msg in msgs {
642        let mut blocks: Vec<LlmContentBlock> = Vec::new();
643        for part in msg.parts.iter() {
644            match part.kind {
645                PartKind::Reasoning => {
646                    let Some(meta) = part.reasoning_meta.as_ref() else {
647                        continue;
648                    };
649                    if meta.is_empty() {
650                        continue;
651                    }
652                    blocks.push(LlmContentBlock::Reasoning {
653                        text: part.content.clone(),
654                        replay: Some(meta.clone()),
655                    });
656                }
657                PartKind::ToolCall => {
658                    let call_id = part.tool_call_id.clone().unwrap_or_default();
659                    let tool_name = part.tool_name.clone().unwrap_or_default();
660                    blocks.push(LlmContentBlock::ToolCall {
661                        call_id,
662                        tool_name,
663                        input_json: part.content.clone(),
664                        replay: part.tool_replay.clone(),
665                    });
666                }
667                PartKind::ToolResult => {
668                    let text = part.render();
669                    let call_id = part.tool_call_id.clone().unwrap_or_default();
670                    blocks.push(LlmContentBlock::ToolResult {
671                        call_id,
672                        content: text,
673                        tool_name: part.tool_name.clone(),
674                    });
675                }
676                _ => {
677                    if let Some(attachment) = attachment_from_part(part)
678                        && matches!(msg.role, MessageRole::User)
679                    {
680                        let attachment_idx = rendered.attachments.len();
681                        rendered.attachments.push(attachment);
682                        blocks.push(LlmContentBlock::Image { attachment_idx });
683                        continue;
684                    }
685
686                    let mut text = render_part_for_chat(msg.role, part);
687                    if text.trim().is_empty() {
688                        continue;
689                    }
690
691                    if matches!(msg.role, MessageRole::System) {
692                        text = format!("Runtime note:\n{text}");
693                    }
694
695                    blocks.push(LlmContentBlock::Text {
696                        text: text.into(),
697                        response_meta: if matches!(part.kind, PartKind::Text | PartKind::Prose) {
698                            part.response_meta.clone()
699                        } else {
700                            None
701                        },
702                        cache_breakpoint: false,
703                    });
704                }
705            }
706        }
707        if blocks.is_empty() {
708            continue;
709        }
710        rendered
711            .messages
712            .push(LlmMessage::new(llm_role_for_message(msg.role), blocks));
713    }
714}
715
716fn llm_role_for_message(role: MessageRole) -> LlmRole {
717    match role {
718        MessageRole::User => LlmRole::User,
719        MessageRole::Assistant => LlmRole::Assistant,
720        MessageRole::System => LlmRole::System,
721    }
722}
723
724#[cfg(test)]
725mod tests {
726    use super::*;
727
728    fn part(kind: PartKind, content: &str) -> Part {
729        Part {
730            id: "p0".to_string(),
731            kind,
732            content: content.to_string(),
733            attachment: None,
734            tool_call_id: None,
735            tool_name: None,
736            tool_replay: None,
737            prune_state: PruneState::Intact,
738            reasoning_meta: None,
739            response_meta: None,
740        }
741    }
742
743    fn test_attachment_ref(byte_len: u64) -> AttachmentRef {
744        AttachmentRef {
745            id: crate::AttachmentId::new("att-test"),
746            media_type: crate::MediaType::Image(crate::ImageMediaType::Png),
747            byte_len,
748            width: None,
749            height: None,
750            label: None,
751        }
752    }
753
754    fn image_part(bytes: &[u8]) -> Part {
755        Part {
756            id: "p0".to_string(),
757            kind: PartKind::Image,
758            content: String::new(),
759            attachment: Some(PartAttachment {
760                reference: test_attachment_ref(bytes.len() as u64),
761            }),
762            tool_call_id: None,
763            tool_name: None,
764            tool_replay: None,
765            prune_state: PruneState::Intact,
766            reasoning_meta: None,
767            response_meta: None,
768        }
769    }
770
771    #[test]
772    fn render_transcript_prompt_orders_turns_oldest_first() {
773        let msgs = vec![
774            Message {
775                id: "m0".to_string(),
776                role: MessageRole::User,
777                parts: vec![part(PartKind::Text, "first")].into(),
778                origin: None,
779            },
780            Message {
781                id: "m1".to_string(),
782                role: MessageRole::Assistant,
783                parts: vec![part(PartKind::Prose, "reply one")].into(),
784                origin: None,
785            },
786            Message {
787                id: "m2".to_string(),
788                role: MessageRole::User,
789                parts: vec![part(PartKind::Text, "second")].into(),
790                origin: None,
791            },
792        ];
793
794        let rendered = render_transcript_prompt(&msgs);
795        let text = block_text(&rendered.messages[0], 0);
796
797        assert!(text.contains("=== Turn 1 ===\nUser:\nfirst"));
798        assert!(text.contains("Assistant (Lash, continuing this transcript):\nreply one"));
799        assert!(text.contains("=== Turn 2 ===\nUser:\nsecond"));
800    }
801
802    fn block_text(msg: &LlmMessage, idx: usize) -> &str {
803        match msg.blocks.get(idx) {
804            Some(LlmContentBlock::Text { text, .. }) => text.as_ref(),
805            Some(other) => panic!("expected Text block, got {other:?}"),
806            None => panic!("missing block at index {idx}"),
807        }
808    }
809
810    #[test]
811    fn render_prompt_repl_preserves_message_boundaries() {
812        let msgs = vec![
813            Message {
814                id: "m1".to_string(),
815                role: MessageRole::User,
816                parts: vec![part(PartKind::Text, "first")].into(),
817                origin: None,
818            },
819            Message {
820                id: "m2".to_string(),
821                role: MessageRole::Assistant,
822                parts: vec![
823                    part(PartKind::Prose, "reply one"),
824                    part(PartKind::Code, "x = 1"),
825                ]
826                .into(),
827                origin: None,
828            },
829            Message {
830                id: "m3".to_string(),
831                role: MessageRole::User,
832                parts: vec![part(PartKind::Text, "second")].into(),
833                origin: None,
834            },
835        ];
836
837        let rendered = render_prompt(&msgs);
838        assert_eq!(rendered.messages.len(), 3);
839        assert_eq!(block_text(&rendered.messages[0], 0), "first");
840        assert!(block_text(&rendered.messages[1], 0).contains("reply one"));
841        assert_eq!(block_text(&rendered.messages[1], 1), "x = 1");
842        assert_eq!(block_text(&rendered.messages[2], 0), "second");
843    }
844
845    #[test]
846    fn render_structured_prompt_preserves_tool_protocol_and_user_images() {
847        let msgs = vec![
848            Message {
849                id: "m0".to_string(),
850                role: MessageRole::System,
851                parts: vec![part(PartKind::Text, "note")].into(),
852                origin: None,
853            },
854            Message {
855                id: "m1".to_string(),
856                role: MessageRole::User,
857                parts: vec![part(PartKind::Text, "show this"), image_part(&[1, 2, 3])].into(),
858                origin: None,
859            },
860            Message {
861                id: "m2".to_string(),
862                role: MessageRole::Assistant,
863                parts: vec![Part {
864                    id: "m2.p0".to_string(),
865                    kind: PartKind::ToolCall,
866                    content: r#"{"path":"README.md"}"#.to_string(),
867                    attachment: None,
868                    tool_call_id: Some("tc1".to_string()),
869                    tool_name: Some("read_file".to_string()),
870                    tool_replay: None,
871                    prune_state: PruneState::Intact,
872                    reasoning_meta: None,
873                    response_meta: None,
874                }]
875                .into(),
876                origin: None,
877            },
878            Message {
879                id: "m3".to_string(),
880                role: MessageRole::User,
881                parts: vec![Part {
882                    id: "m3.p0".to_string(),
883                    kind: PartKind::ToolResult,
884                    content: "ok".to_string(),
885                    attachment: None,
886                    tool_call_id: Some("tc1".to_string()),
887                    tool_name: Some("read_file".to_string()),
888                    tool_replay: None,
889                    prune_state: PruneState::Intact,
890                    reasoning_meta: None,
891                    response_meta: None,
892                }]
893                .into(),
894                origin: None,
895            },
896        ];
897
898        let rendered = render_structured_prompt(&msgs);
899        assert_eq!(rendered.messages.len(), 4);
900        assert_eq!(rendered.messages[0].role, LlmRole::System);
901        assert_eq!(block_text(&rendered.messages[0], 0), "Runtime note:\nnote");
902        // User message has text + image blocks bundled together.
903        assert_eq!(rendered.messages[1].role, LlmRole::User);
904        assert!(matches!(
905            rendered.messages[1].blocks[0],
906            LlmContentBlock::Text { .. }
907        ));
908        assert!(matches!(
909            rendered.messages[1].blocks[1],
910            LlmContentBlock::Image { attachment_idx: 0 }
911        ));
912        assert_eq!(rendered.attachments.len(), 1);
913        assert!(matches!(
914            rendered.messages[2].blocks[0],
915            LlmContentBlock::ToolCall { .. }
916        ));
917        assert!(matches!(
918            rendered.messages[3].blocks[0],
919            LlmContentBlock::ToolResult { .. }
920        ));
921    }
922
923    #[test]
924    fn render_structured_prompt_preserves_empty_tool_results() {
925        let msgs = vec![
926            Message {
927                id: "m0".to_string(),
928                role: MessageRole::Assistant,
929                parts: vec![Part {
930                    id: "m0.p0".to_string(),
931                    kind: PartKind::ToolCall,
932                    content: r#"{"question":"Pick one"}"#.to_string(),
933                    attachment: None,
934                    tool_call_id: Some("ask_1".to_string()),
935                    tool_name: Some("ask".to_string()),
936                    tool_replay: None,
937                    prune_state: PruneState::Intact,
938                    reasoning_meta: None,
939                    response_meta: None,
940                }]
941                .into(),
942                origin: None,
943            },
944            Message {
945                id: "m1".to_string(),
946                role: MessageRole::User,
947                parts: vec![Part {
948                    id: "m1.p0".to_string(),
949                    kind: PartKind::ToolResult,
950                    content: String::new(),
951                    attachment: None,
952                    tool_call_id: Some("ask_1".to_string()),
953                    tool_name: Some("ask".to_string()),
954                    tool_replay: None,
955                    prune_state: PruneState::Intact,
956                    reasoning_meta: None,
957                    response_meta: None,
958                }]
959                .into(),
960                origin: None,
961            },
962        ];
963
964        let rendered = render_structured_prompt(&msgs);
965        assert_eq!(rendered.messages.len(), 2);
966        match &rendered.messages[0].blocks[0] {
967            LlmContentBlock::ToolCall {
968                call_id, tool_name, ..
969            } => {
970                assert_eq!(call_id, "ask_1");
971                assert_eq!(tool_name, "ask");
972            }
973            other => panic!("expected ToolCall, got {other:?}"),
974        }
975        match &rendered.messages[1].blocks[0] {
976            LlmContentBlock::ToolResult {
977                call_id, content, ..
978            } => {
979                assert_eq!(call_id, "ask_1");
980                assert!(content.is_empty());
981            }
982            other => panic!("expected ToolResult, got {other:?}"),
983        }
984    }
985
986    #[test]
987    fn render_transcript_prompt_collects_images() {
988        let msgs = vec![Message {
989            id: "m0".to_string(),
990            role: MessageRole::User,
991            parts: vec![image_part(&[9, 8, 7])].into(),
992            origin: None,
993        }];
994
995        let rendered = render_transcript_prompt(&msgs);
996        let text = block_text(&rendered.messages[0], 0);
997        assert!(text.contains("[Image attached]"));
998        assert_eq!(rendered.attachments.len(), 1);
999    }
1000
1001    #[test]
1002    fn render_transcript_prompt_omits_missing_assistant_placeholder_for_current_turn() {
1003        let msgs = vec![
1004            Message {
1005                id: "m0".to_string(),
1006                role: MessageRole::User,
1007                parts: vec![part(PartKind::Text, "first")].into(),
1008                origin: None,
1009            },
1010            Message {
1011                id: "m1".to_string(),
1012                role: MessageRole::Assistant,
1013                parts: vec![part(PartKind::Prose, "reply one")].into(),
1014                origin: None,
1015            },
1016            Message {
1017                id: "m2".to_string(),
1018                role: MessageRole::User,
1019                parts: vec![part(PartKind::Text, "second")].into(),
1020                origin: None,
1021            },
1022        ];
1023
1024        let rendered = render_transcript_prompt(&msgs);
1025        let text = block_text(&rendered.messages[0], 0);
1026
1027        assert!(text.contains("=== Turn 2 ===\nUser:\nsecond"));
1028        assert!(!text.contains("=== Turn 2 ===\nUser:\nsecond\n\nAssistant (Lash, continuing this transcript):\n[No assistant content recorded]"));
1029    }
1030
1031    #[test]
1032    fn render_transcript_prompt_preserves_tool_name_for_assistant_tool_calls() {
1033        let msgs = vec![
1034            Message {
1035                id: "m0".to_string(),
1036                role: MessageRole::User,
1037                parts: vec![part(PartKind::Text, "what time is it")].into(),
1038                origin: None,
1039            },
1040            Message {
1041                id: "m1".to_string(),
1042                role: MessageRole::Assistant,
1043                parts: vec![Part {
1044                    id: "m1.p0".to_string(),
1045                    kind: PartKind::ToolCall,
1046                    content: r#"{"cmd":"date"}"#.to_string(),
1047                    attachment: None,
1048                    tool_call_id: Some("tc1".to_string()),
1049                    tool_name: Some("exec_command".to_string()),
1050                    tool_replay: None,
1051                    prune_state: PruneState::Intact,
1052                    reasoning_meta: None,
1053                    response_meta: None,
1054                }]
1055                .into(),
1056                origin: None,
1057            },
1058        ];
1059
1060        let rendered = render_transcript_prompt(&msgs);
1061        let text = block_text(&rendered.messages[0], 0);
1062
1063        assert!(text.contains(r#"exec_command({"cmd":"date"})"#));
1064    }
1065
1066    #[test]
1067    fn render_transcript_prompt_omits_runtime_notes_section() {
1068        let msgs = vec![Message {
1069            id: "m0".to_string(),
1070            role: MessageRole::User,
1071            parts: vec![part(PartKind::Text, "hi")].into(),
1072            origin: None,
1073        }];
1074
1075        let rendered = render_transcript_prompt(&msgs);
1076        let text = block_text(&rendered.messages[0], 0);
1077        assert!(!text.contains("Runtime Notes:"));
1078    }
1079
1080    #[test]
1081    fn prompt_resume_safety_accepts_completed_tool_history() {
1082        let msgs = vec![
1083            Message {
1084                id: "m0".to_string(),
1085                role: MessageRole::Assistant,
1086                parts: vec![Part {
1087                    id: "m0.p0".to_string(),
1088                    kind: PartKind::ToolCall,
1089                    content: r#"{"path":"README.md"}"#.to_string(),
1090                    attachment: None,
1091                    tool_call_id: Some("tc1".to_string()),
1092                    tool_name: Some("read_file".to_string()),
1093                    tool_replay: None,
1094                    prune_state: PruneState::Intact,
1095                    reasoning_meta: None,
1096                    response_meta: None,
1097                }]
1098                .into(),
1099                origin: None,
1100            },
1101            Message {
1102                id: "m1".to_string(),
1103                role: MessageRole::User,
1104                parts: vec![Part {
1105                    id: "m1.p0".to_string(),
1106                    kind: PartKind::ToolResult,
1107                    content: "ok".to_string(),
1108                    attachment: None,
1109                    tool_call_id: Some("tc1".to_string()),
1110                    tool_name: Some("read_file".to_string()),
1111                    tool_replay: None,
1112                    prune_state: PruneState::Intact,
1113                    reasoning_meta: None,
1114                    response_meta: None,
1115                }]
1116                .into(),
1117                origin: None,
1118            },
1119        ];
1120
1121        assert!(messages_are_prompt_resume_safe(&msgs));
1122    }
1123
1124    #[test]
1125    fn reasoning_parts_survive_snapshot_but_never_reach_the_model() {
1126        let reasoning_part = Part {
1127            id: "m1.p0".to_string(),
1128            kind: PartKind::Reasoning,
1129            content: "Thinking about how to answer.".to_string(),
1130            attachment: None,
1131            tool_call_id: None,
1132            tool_name: None,
1133            tool_replay: None,
1134            prune_state: PruneState::Intact,
1135            reasoning_meta: None,
1136            response_meta: None,
1137        };
1138
1139        let msgs = vec![Message {
1140            id: "m1".to_string(),
1141            role: MessageRole::Assistant,
1142            parts: vec![
1143                reasoning_part.clone(),
1144                part(PartKind::Prose, "Here is the answer."),
1145            ]
1146            .into(),
1147            origin: None,
1148        }];
1149
1150        // JSON round-trip preserves the reasoning part — the snapshot
1151        // layer must not silently drop it, otherwise replays would lose
1152        // the trace.
1153        let serialized = serde_json::to_string(&msgs).expect("serialize messages");
1154        let deserialized: Vec<Message> =
1155            serde_json::from_str(&serialized).expect("deserialize messages");
1156        assert_eq!(deserialized[0].parts.len(), 2);
1157        assert!(matches!(deserialized[0].parts[0].kind, PartKind::Reasoning));
1158        assert_eq!(
1159            deserialized[0].parts[0].content,
1160            "Thinking about how to answer."
1161        );
1162
1163        // But the rendered LLM prompt must NOT include the reasoning
1164        // content in any assistant TEXT block — reasoning travels as its
1165        // own block kind so adapters that don't understand it can drop
1166        // without corrupting the visible transcript.
1167        let rendered = render_structured_prompt(&msgs);
1168        assert_eq!(rendered.messages.len(), 1);
1169        assert_eq!(rendered.messages[0].role, LlmRole::Assistant);
1170        // Without `reasoning_meta`, the reasoning part is dropped entirely,
1171        // so the assistant turn contains only the prose block.
1172        assert_eq!(rendered.messages[0].blocks.len(), 1);
1173        assert!(matches!(
1174            &rendered.messages[0].blocks[0],
1175            LlmContentBlock::Text { text, .. } if text.as_ref() == "Here is the answer."
1176        ));
1177
1178        // When the assistant message consists solely of a display-only
1179        // reasoning part (no encrypted payload), no message is sent at
1180        // all.
1181        let reasoning_only = vec![Message {
1182            id: "m2".to_string(),
1183            role: MessageRole::Assistant,
1184            parts: vec![reasoning_part].into(),
1185            origin: None,
1186        }];
1187        let rendered_only = render_structured_prompt(&reasoning_only);
1188        assert!(rendered_only.messages.is_empty());
1189    }
1190
1191    #[test]
1192    fn prompt_resume_safety_rejects_unmatched_tool_calls() {
1193        let msgs = vec![Message {
1194            id: "m0".to_string(),
1195            role: MessageRole::Assistant,
1196            parts: vec![Part {
1197                id: "m0.p0".to_string(),
1198                kind: PartKind::ToolCall,
1199                content: r#"{"path":"README.md"}"#.to_string(),
1200                attachment: None,
1201                tool_call_id: Some("tc1".to_string()),
1202                tool_name: Some("read_file".to_string()),
1203                tool_replay: None,
1204                prune_state: PruneState::Intact,
1205                reasoning_meta: None,
1206                response_meta: None,
1207            }]
1208            .into(),
1209            origin: None,
1210        }];
1211
1212        assert!(!messages_are_prompt_resume_safe(&msgs));
1213    }
1214
1215    // ─── Reasoning-part roundtrip (fix 1.3b) ──────────────────────────
1216    //
1217    // Provider reasoning items can carry replay metadata that the adapter
1218    // re-emits on the next turn. The session-model layer stores these parts
1219    // so they survive resume/snapshot and flows them through as
1220    // `kind == "reasoning"` LlmMessages.
1221
1222    fn reasoning_part_fixture(encrypted: Option<&str>) -> Part {
1223        Part {
1224            id: "m0.p0".to_string(),
1225            kind: PartKind::Reasoning,
1226            content: "Thinking.".to_string(),
1227            attachment: None,
1228            tool_call_id: None,
1229            tool_name: None,
1230            tool_replay: None,
1231            prune_state: PruneState::Intact,
1232            reasoning_meta: encrypted.map(|encrypted| ProviderReasoningReplay {
1233                item_id: Some("rs_xyz".to_string()),
1234                summary: vec!["Thinking.".to_string()],
1235                encrypted_content: Some(encrypted.to_string()),
1236                signature: None,
1237                redacted: false,
1238            }),
1239            response_meta: None,
1240        }
1241    }
1242
1243    #[test]
1244    fn reasoning_part_roundtrips_through_snapshot_serde() {
1245        let msgs = vec![Message {
1246            id: "m0".to_string(),
1247            role: MessageRole::Assistant,
1248            parts: vec![reasoning_part_fixture(Some("CIPHER=="))].into(),
1249            origin: None,
1250        }];
1251        let serialized = serde_json::to_string(&msgs).expect("serialize");
1252        let deserialized: Vec<Message> = serde_json::from_str(&serialized).expect("deserialize");
1253        assert_eq!(deserialized[0].parts.len(), 1);
1254        let part = &deserialized[0].parts[0];
1255        assert!(matches!(part.kind, PartKind::Reasoning));
1256        let meta = part.reasoning_meta.as_ref().expect("meta survives");
1257        assert_eq!(meta.item_id.as_deref(), Some("rs_xyz"));
1258        assert_eq!(meta.summary, vec!["Thinking.".to_string()]);
1259        assert_eq!(meta.encrypted_content.as_deref(), Some("CIPHER=="));
1260    }
1261
1262    #[test]
1263    fn reasoning_part_roundtrips_when_snapshot_predates_field() {
1264        // Older snapshots written before fix 1.3b have no
1265        // `reasoning_meta` column. The field must default to `None`
1266        // and the deserializer must accept the legacy shape.
1267        let legacy = r#"[{
1268            "id":"m0","role":"Assistant",
1269            "parts":[{
1270                "id":"m0.p0","kind":"Prose","content":"Hi",
1271                "prune_state":"Intact"
1272            }]
1273        }]"#;
1274        let msgs: Vec<Message> = serde_json::from_str(legacy).expect("legacy snapshot");
1275        assert!(msgs[0].parts[0].reasoning_meta.is_none());
1276    }
1277
1278    #[test]
1279    fn reasoning_parts_never_flow_to_rendered_prompt_as_text() {
1280        // Whether or not the reasoning item carries an encrypted blob,
1281        // it must NEVER be flattened into assistant text content.
1282        // Without an encrypted blob the adapter also drops it entirely
1283        // (no point re-feeding a display-only summary).
1284        let display_only = vec![Message {
1285            id: "m0".to_string(),
1286            role: MessageRole::Assistant,
1287            parts: vec![reasoning_part_fixture(None)].into(),
1288            origin: None,
1289        }];
1290        let rendered = render_structured_prompt(&display_only);
1291        assert!(
1292            rendered.messages.is_empty(),
1293            "display-only reasoning must not reach the prompt"
1294        );
1295
1296        // With encrypted content, a single Reasoning block is emitted
1297        // that adapters can re-emit via their native reasoning channel.
1298        let replayable = vec![Message {
1299            id: "m0".to_string(),
1300            role: MessageRole::Assistant,
1301            parts: vec![reasoning_part_fixture(Some("CIPHER=="))].into(),
1302            origin: None,
1303        }];
1304        let rendered = render_structured_prompt(&replayable);
1305        assert_eq!(rendered.messages.len(), 1);
1306        match &rendered.messages[0].blocks[0] {
1307            LlmContentBlock::Reasoning { replay, .. } => {
1308                let replay = replay.as_ref().expect("reasoning replay");
1309                assert_eq!(replay.encrypted_content.as_deref(), Some("CIPHER=="));
1310                assert_eq!(replay.item_id.as_deref(), Some("rs_xyz"));
1311                assert_eq!(replay.summary, vec!["Thinking.".to_string()]);
1312            }
1313            other => panic!("expected Reasoning block, got {other:?}"),
1314        }
1315        // Sanity: transcript rendering never includes reasoning text.
1316        let transcript = render_transcript_prompt(&replayable);
1317        let transcript_text = block_text(&transcript.messages[0], 0);
1318        assert!(!transcript_text.contains("Thinking."));
1319        assert!(!transcript_text.contains("CIPHER=="));
1320    }
1321
1322    #[test]
1323    fn reasoning_parts_are_zero_for_prune_accounting() {
1324        // The rolling-history plugin's prune logic is driven by
1325        // `prompt_char_count`. Reasoning parts are not user-visible,
1326        // so they must not count against the prompt budget.
1327        let part = reasoning_part_fixture(Some("X=="));
1328        assert_eq!(part.prompt_char_count(), 0);
1329    }
1330}
lash_sansio/session_model/message.rs

lash_sansio/session_model/
message.rs
