lash_sansio/session_model/
mod.rs

1pub mod message;
2pub mod prompt;
3
4pub use message::{
5    BaseRenderCache, Message, MessageRole, MessageSequence, Part, PartAttachment, PartKind,
6    PruneState, RenderedPrompt, append_rendered_prompt, messages_are_prompt_resume_safe,
7    render_prompt, render_transcript_prompt, shared_parts,
8};
9pub use prompt::{
10    MAIN_AGENT_INTRO, PromptBuiltin, PromptLayer, PromptSlot, PromptSlotLayer, PromptTemplate,
11    PromptTemplateEntry, PromptTemplateSection, ResolvedPromptLayer, default_prompt_template,
12    resolve_prompt_layers,
13};
14
15use std::sync::Arc;
16
17use crate::MessageOrigin;
18use crate::ToolDefinition;
19use crate::llm::types::LlmToolSpec;
20use crate::plugin::{CheckpointKind, PluginMessage, PluginRuntimeEvent};
21
22/// Durable protocol payload stored in session history.
23#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
24pub struct ProtocolEvent {
25    pub plugin_id: String,
26    pub payload: serde_json::Value,
27}
28
29impl ProtocolEvent {
30    pub fn typed<T>(plugin_id: impl Into<String>, event: T) -> Result<Self, serde_json::Error>
31    where
32        T: serde::Serialize,
33    {
34        Ok(Self {
35            plugin_id: plugin_id.into(),
36            payload: serde_json::to_value(event)?,
37        })
38    }
39
40    pub fn decode<T>(&self, expected_plugin_id: &str) -> Result<Option<T>, serde_json::Error>
41    where
42        T: for<'de> serde::Deserialize<'de>,
43    {
44        if self.plugin_id != expected_plugin_id {
45            return Ok(None);
46        }
47        serde_json::from_value(self.payload.clone()).map(Some)
48    }
49}
50
51/// Typed node accepted at session-graph append boundaries.
52#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
53#[serde(tag = "kind", rename_all = "snake_case")]
54// justification: append nodes are public durable DTOs kept inline to preserve their established Rust construction API.
55#[allow(clippy::large_enum_variant)]
56pub enum SessionAppendNode {
57    Message {
58        message: PluginMessage,
59        #[serde(default, skip_serializing_if = "Option::is_none")]
60        caused_by: Option<crate::CausalRef>,
61    },
62    ProtocolEvent {
63        event: ProtocolEvent,
64        #[serde(default, skip_serializing_if = "Option::is_none")]
65        caused_by: Option<crate::CausalRef>,
66    },
67    Plugin {
68        plugin_type: String,
69        #[serde(default)]
70        body: serde_json::Value,
71        #[serde(default, skip_serializing_if = "Option::is_none")]
72        caused_by: Option<crate::CausalRef>,
73    },
74}
75
76impl SessionAppendNode {
77    pub fn message(message: PluginMessage) -> Self {
78        Self::Message {
79            message,
80            caused_by: None,
81        }
82    }
83
84    pub fn plugin(plugin_type: impl Into<String>, body: serde_json::Value) -> Self {
85        Self::Plugin {
86            plugin_type: plugin_type.into(),
87            body,
88            caused_by: None,
89        }
90    }
91
92    pub fn protocol_event(event: ProtocolEvent) -> Self {
93        Self::ProtocolEvent {
94            event,
95            caused_by: None,
96        }
97    }
98
99    pub fn with_caused_by(mut self, caused_by: crate::CausalRef) -> Self {
100        match &mut self {
101            Self::Message {
102                caused_by: cause, ..
103            }
104            | Self::ProtocolEvent {
105                caused_by: cause, ..
106            }
107            | Self::Plugin {
108                caused_by: cause, ..
109            } => *cause = Some(caused_by),
110        }
111        self
112    }
113}
114
115/// Durable semantic history stored in the session graph and replayed into
116/// future prompts. Unlike [`SessionStreamEvent`], these records are committed
117/// state rather than transient UI/progress signals.
118#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
119// justification: the generic protocol event is caller-defined durable state and must retain the public inline history shape.
120#[allow(clippy::large_enum_variant)]
121pub enum SessionHistoryRecord<PE = ()> {
122    Conversation(ConversationRecord),
123    Protocol(PE),
124}
125
126#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
127pub struct ConversationRecord {
128    pub id: String,
129    pub role: MessageRole,
130    pub parts: Arc<Vec<Part>>,
131    #[serde(default, skip_serializing_if = "Option::is_none")]
132    pub origin: Option<MessageOrigin>,
133}
134
135#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
136pub struct AcceptedInjectedTurnInput {
137    #[serde(default, skip_serializing_if = "Option::is_none")]
138    pub id: Option<String>,
139    pub message: PluginMessage,
140}
141
142impl ConversationRecord {
143    pub fn from_message(message: Message) -> Self {
144        Self {
145            id: message.id,
146            role: message.role,
147            parts: message.parts,
148            origin: message.origin,
149        }
150    }
151
152    pub fn to_message(&self) -> Message {
153        Message {
154            id: self.id.clone(),
155            role: self.role,
156            parts: Arc::clone(&self.parts),
157            origin: self.origin.clone(),
158        }
159    }
160}
161
162/// Token usage statistics from an LLM call.
163#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
164pub struct TokenUsage {
165    pub input_tokens: i64,
166    pub output_tokens: i64,
167    pub cache_read_input_tokens: i64,
168    pub cache_write_input_tokens: i64,
169    pub reasoning_output_tokens: i64,
170}
171
172impl TokenUsage {
173    pub fn total(&self) -> i64 {
174        self.input_tokens
175            + self.output_tokens
176            + self.cache_read_input_tokens
177            + self.cache_write_input_tokens
178    }
179
180    pub fn input_total(&self) -> i64 {
181        self.input_tokens + self.cache_read_input_tokens + self.cache_write_input_tokens
182    }
183
184    pub fn add(&mut self, other: &TokenUsage) {
185        self.input_tokens += other.input_tokens;
186        self.output_tokens += other.output_tokens;
187        self.cache_read_input_tokens += other.cache_read_input_tokens;
188        self.cache_write_input_tokens += other.cache_write_input_tokens;
189        self.reasoning_output_tokens += other.reasoning_output_tokens;
190    }
191}
192
193/// Structured error payload carried on [`SessionStreamEvent::Error`] (and
194/// [`SessionStreamEvent::RetryStatus`]).
195///
196/// Durability: this type appears inside persisted session snapshots and turn
197/// checkpoints, so every field added after the initial shape must stay
198/// additive — `#[serde(default)]` on decode and
199/// `#[serde(skip_serializing_if = "Option::is_none")]` on encode — to keep
200/// old snapshots decodable and new snapshots readable by older readers.
201/// Transient runtime-stream signal for live consumers.
202///
203/// These events may be partial, duplicated, or display-only and are not proof
204/// of durable session history. Persisted semantic history uses
205/// [`SessionHistoryRecord`].
206#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
207pub struct ErrorEnvelope {
208    pub kind: String,
209    #[serde(default, skip_serializing_if = "Option::is_none")]
210    pub code: Option<String>,
211    #[serde(default, skip_serializing_if = "Option::is_none")]
212    pub terminal_reason: Option<crate::llm::types::LlmTerminalReason>,
213    pub user_message: String,
214    #[serde(default, skip_serializing_if = "Option::is_none")]
215    pub raw: Option<String>,
216    /// Whether the failing operation is safe to retry, when the source
217    /// carries a typed signal (provider transports classify retryability).
218    /// `None` means the source did not know.
219    #[serde(default, skip_serializing_if = "Option::is_none")]
220    pub retryable: Option<bool>,
221    /// Typed provider-failure classification, set only when the error came
222    /// from an LLM provider/transport failure whose kind was classified
223    /// (an unclassified `Unknown` kind is surfaced as `None`).
224    #[serde(default, skip_serializing_if = "Option::is_none")]
225    pub provider_failure_kind: Option<crate::llm::types::ProviderFailureKind>,
226}
227
228#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
229#[serde(tag = "type")]
230// justification: this public streaming DTO stays inline to avoid per-event allocation and preserve consumer pattern matching.
231#[allow(clippy::large_enum_variant)]
232pub enum SessionStreamEvent {
233    #[serde(rename = "text_delta")]
234    TextDelta { content: String },
235    /// Streaming update for the model's reasoning summary ("thinking").
236    /// The UI renders these incrementally in a muted/italic style;
237    /// reasoning content is never fed back to the model on subsequent
238    /// turns.
239    #[serde(rename = "reasoning_delta")]
240    ReasoningDelta { content: String },
241    #[serde(rename = "tool_call")]
242    ToolCall {
243        #[serde(default, skip_serializing_if = "Option::is_none")]
244        call_id: Option<String>,
245        name: String,
246        args: serde_json::Value,
247        output: crate::ToolCallOutput,
248        duration_ms: u64,
249    },
250    #[serde(rename = "tool_call_start")]
251    ToolCallStart {
252        #[serde(default, skip_serializing_if = "Option::is_none")]
253        call_id: Option<String>,
254        name: String,
255        args: serde_json::Value,
256    },
257    #[serde(rename = "message")]
258    Message { text: String, kind: String },
259    #[serde(rename = "llm_request")]
260    LlmRequest {
261        protocol_iteration: usize,
262        message_count: usize,
263        tool_list: String,
264    },
265    #[serde(rename = "llm_response")]
266    LlmResponse {
267        protocol_iteration: usize,
268        content: String,
269        duration_ms: u64,
270    },
271    #[serde(rename = "token_usage")]
272    TokenUsage {
273        protocol_iteration: usize,
274        usage: TokenUsage,
275        cumulative: TokenUsage,
276    },
277    #[serde(rename = "child_token_usage")]
278    ChildTokenUsage {
279        session_id: String,
280        source: String,
281        model: String,
282        protocol_iteration: usize,
283        usage: TokenUsage,
284        cumulative: TokenUsage,
285    },
286    #[serde(rename = "retry_status")]
287    RetryStatus {
288        wait_seconds: u64,
289        attempt: usize,
290        max_attempts: usize,
291        reason: String,
292        #[serde(default, skip_serializing_if = "Option::is_none")]
293        envelope: Option<ErrorEnvelope>,
294    },
295    #[serde(rename = "injected_turn_input_accepted")]
296    InjectedTurnInputAccepted {
297        inputs: Vec<AcceptedInjectedTurnInput>,
298        checkpoint: CheckpointKind,
299    },
300    #[serde(rename = "injected_messages_committed")]
301    InjectedMessagesCommitted {
302        messages: Vec<PluginMessage>,
303        checkpoint: CheckpointKind,
304    },
305    #[serde(rename = "plugin_event")]
306    PluginEvent {
307        plugin_id: String,
308        event: PluginRuntimeEvent,
309    },
310    /// Semantic result for a completed turn. `Done` remains the machine
311    /// lifecycle marker emitted after this event.
312    #[serde(rename = "turn_outcome")]
313    TurnOutcome { outcome: TurnOutcome },
314    #[serde(rename = "done")]
315    Done,
316    #[serde(rename = "error")]
317    Error {
318        message: String,
319        #[serde(default, skip_serializing_if = "Option::is_none")]
320        envelope: Option<ErrorEnvelope>,
321    },
322}
323
324#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
325#[serde(rename_all = "snake_case")]
326pub enum TurnOutcome {
327    Finished(TurnFinish),
328    AgentFrameSwitch {
329        frame_id: String,
330        task: String,
331        #[serde(default, skip_serializing_if = "Vec::is_empty")]
332        initial_nodes: Vec<SessionAppendNode>,
333    },
334    Stopped(TurnStop),
335}
336
337#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
338#[serde(rename_all = "snake_case")]
339pub enum TurnFinish {
340    AssistantMessage {
341        text: String,
342    },
343    FinalValue {
344        value: serde_json::Value,
345    },
346    ToolValue {
347        tool_name: String,
348        value: serde_json::Value,
349    },
350}
351
352#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
353#[serde(rename_all = "snake_case")]
354pub enum TurnStop {
355    Cancelled,
356    Incomplete,
357    InvalidInput,
358    MaxTurns,
359    ToolFailure,
360    ProviderError,
361    PluginAbort,
362    RuntimeError,
363    SubmittedError {
364        value: serde_json::Value,
365    },
366    ToolError {
367        tool_name: String,
368        value: serde_json::Value,
369    },
370}
371
372#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
373pub struct TurnTerminationPolicyState {
374    turn_limit_final_scheduled: bool,
375}
376
377impl Default for TurnTerminationPolicyState {
378    fn default() -> Self {
379        Self::new()
380    }
381}
382
383impl TurnTerminationPolicyState {
384    pub fn new() -> Self {
385        Self {
386            turn_limit_final_scheduled: false,
387        }
388    }
389
390    pub fn should_force_exit_after_grace_turn(&self) -> bool {
391        self.turn_limit_final_scheduled
392    }
393
394    pub fn turn_limit_final_to_schedule(
395        &self,
396        protocol_iteration: usize,
397        protocol_run_offset: usize,
398        max_turns: Option<usize>,
399    ) -> Option<usize> {
400        if self.turn_limit_final_scheduled {
401            return None;
402        }
403        let max = max_turns?;
404        if protocol_iteration < protocol_run_offset + max {
405            return None;
406        }
407        Some(max)
408    }
409
410    pub fn mark_turn_limit_final_scheduled(&mut self) {
411        self.turn_limit_final_scheduled = true;
412    }
413}
414
415pub fn make_error_envelope(
416    kind: &str,
417    code: Option<&str>,
418    terminal_reason: Option<crate::llm::types::LlmTerminalReason>,
419    user_message: impl Into<String>,
420    raw: Option<String>,
421) -> ErrorEnvelope {
422    let user_message = user_message.into();
423    ErrorEnvelope {
424        kind: kind.to_string(),
425        code: code.map(str::to_string),
426        terminal_reason,
427        user_message,
428        raw: raw.map(|s| truncate_raw_error(s.trim())),
429        retryable: None,
430        provider_failure_kind: None,
431    }
432}
433
434pub fn make_error_event(
435    kind: &str,
436    code: Option<&str>,
437    user_message: impl Into<String>,
438    raw: Option<String>,
439) -> SessionStreamEvent {
440    let user_message = user_message.into();
441    SessionStreamEvent::Error {
442        message: user_message.clone(),
443        envelope: Some(make_error_envelope(kind, code, None, user_message, raw)),
444    }
445}
446
447pub fn truncate_raw_error(s: &str) -> String {
448    const MAX_RAW: usize = 4000;
449    let raw_len = s.chars().count();
450    if raw_len <= MAX_RAW {
451        return s.to_string();
452    }
453    let keep = MAX_RAW / 2;
454    let head = s.chars().take(keep).collect::<String>();
455    let tail = s
456        .chars()
457        .rev()
458        .take(keep)
459        .collect::<Vec<_>>()
460        .into_iter()
461        .rev()
462        .collect::<String>();
463    let omitted = raw_len.saturating_sub(keep * 2);
464    format!("{head}\n\n... ({omitted} chars omitted) ...\n\n{tail}")
465}
466
467pub fn reassign_part_ids(message_id: &str, parts: &mut [Part]) {
468    for (idx, part) in parts.iter_mut().enumerate() {
469        part.id = format!("{message_id}.p{idx}");
470    }
471}
472
473pub fn model_tool_specs_iter<'a>(
474    tools: impl IntoIterator<Item = &'a ToolDefinition>,
475) -> Vec<LlmToolSpec> {
476    tools
477        .into_iter()
478        .map(|tool| {
479            let model_tool = tool.model_tool();
480            LlmToolSpec {
481                name: model_tool.name,
482                description: model_tool.description,
483                input_schema: model_tool.input_schema,
484                output_schema: model_tool.output_schema,
485            }
486        })
487        .collect()
488}
489
490pub fn model_tool_specs(tools: &[ToolDefinition]) -> Vec<LlmToolSpec> {
491    model_tool_specs_iter(tools.iter())
492}
493
494#[cfg(test)]
495mod tests {
496    use super::{ErrorEnvelope, SessionStreamEvent, TurnOutcome};
497    use crate::llm::types::{LlmTerminalReason, ProviderFailureKind};
498
499    // ─── ErrorEnvelope durable-snapshot compatibility ──────────────────
500    //
501    // `ErrorEnvelope` is persisted inside session snapshots and turn
502    // checkpoints. The retryability fields added after the initial shape
503    // must decode from legacy JSON (absent fields → `None`) and must not
504    // appear on the wire when unset, so old readers keep decoding new
505    // snapshots too.
506
507    #[test]
508    fn error_envelope_decodes_legacy_snapshot_without_retryability_fields() {
509        let legacy = r#"{
510            "kind":"llm_provider",
511            "code":"429",
512            "terminal_reason":"provider_error",
513            "user_message":"LLM error: rate limited",
514            "raw":"{\"error\":\"rate_limited\"}"
515        }"#;
516        let envelope: ErrorEnvelope = serde_json::from_str(legacy).expect("legacy envelope");
517        assert_eq!(envelope.kind, "llm_provider");
518        assert_eq!(envelope.retryable, None);
519        assert_eq!(envelope.provider_failure_kind, None);
520
521        // The legacy shape embedded in a persisted `SessionStreamEvent::Error`
522        // record decodes the same way.
523        let legacy_event = r#"{
524            "type":"error",
525            "message":"LLM error: rate limited",
526            "envelope":{"kind":"llm_provider","user_message":"LLM error: rate limited"}
527        }"#;
528        let event: SessionStreamEvent = serde_json::from_str(legacy_event).expect("legacy event");
529        match event {
530            SessionStreamEvent::Error { envelope, .. } => {
531                let envelope = envelope.expect("envelope");
532                assert_eq!(envelope.retryable, None);
533                assert_eq!(envelope.provider_failure_kind, None);
534            }
535            other => panic!("expected error event, got {other:?}"),
536        }
537    }
538
539    #[test]
540    fn error_envelope_roundtrips_retryability_fields() {
541        let envelope = ErrorEnvelope {
542            kind: "llm_provider".to_string(),
543            code: Some("429".to_string()),
544            terminal_reason: Some(LlmTerminalReason::ProviderError),
545            user_message: "LLM error: rate limited".to_string(),
546            raw: None,
547            retryable: Some(true),
548            provider_failure_kind: Some(ProviderFailureKind::Quota),
549        };
550        let json = serde_json::to_value(&envelope).expect("serialize envelope");
551        assert_eq!(json["retryable"], serde_json::json!(true));
552        assert_eq!(json["provider_failure_kind"], serde_json::json!("quota"));
553        let decoded: ErrorEnvelope = serde_json::from_value(json).expect("decode envelope");
554        assert_eq!(decoded.retryable, Some(true));
555        assert_eq!(
556            decoded.provider_failure_kind,
557            Some(ProviderFailureKind::Quota)
558        );
559    }
560
561    #[test]
562    fn error_envelope_omits_unset_retryability_fields_on_the_wire() {
563        let envelope = ErrorEnvelope {
564            kind: "plugin".to_string(),
565            code: Some("plugin_abort".to_string()),
566            terminal_reason: None,
567            user_message: "stopped".to_string(),
568            raw: None,
569            retryable: None,
570            provider_failure_kind: None,
571        };
572        let json = serde_json::to_value(&envelope).expect("serialize envelope");
573        let object = json.as_object().expect("object");
574        assert!(!object.contains_key("retryable"));
575        assert!(!object.contains_key("provider_failure_kind"));
576    }
577
578    #[test]
579    fn provider_failure_kind_decodes_unknown_future_codes() {
580        // Forward compatibility: a snapshot written by a newer runtime with a
581        // kind this build does not know decodes as `Unknown`.
582        let decoded: ProviderFailureKind =
583            serde_json::from_value(serde_json::json!("some_future_kind")).expect("future kind");
584        assert_eq!(decoded, ProviderFailureKind::Unknown);
585        for kind in [
586            ProviderFailureKind::Transport,
587            ProviderFailureKind::Timeout,
588            ProviderFailureKind::Http,
589            ProviderFailureKind::Stream,
590            ProviderFailureKind::Auth,
591            ProviderFailureKind::Validation,
592            ProviderFailureKind::Quota,
593            ProviderFailureKind::Unsupported,
594            ProviderFailureKind::Unknown,
595        ] {
596            let json = serde_json::to_value(kind).expect("serialize kind");
597            assert_eq!(json, serde_json::json!(kind.code()));
598            let round: ProviderFailureKind = serde_json::from_value(json).expect("decode kind");
599            assert_eq!(round, kind);
600        }
601    }
602
603    #[test]
604    fn agent_frame_switch_decodes_legacy_event_without_initial_nodes() {
605        let legacy = r#"{
606            "type":"turn_outcome",
607            "outcome":{
608                "agent_frame_switch":{
609                    "frame_id":"frame-2",
610                    "task":"continue"
611                }
612            }
613        }"#;
614        let event: SessionStreamEvent =
615            serde_json::from_str(legacy).expect("legacy frame switch event");
616        match event {
617            SessionStreamEvent::TurnOutcome {
618                outcome:
619                    TurnOutcome::AgentFrameSwitch {
620                        frame_id,
621                        task,
622                        initial_nodes,
623                    },
624            } => {
625                assert_eq!(frame_id, "frame-2");
626                assert_eq!(task, "continue");
627                assert!(initial_nodes.is_empty());
628            }
629            other => panic!("expected agent-frame switch event, got {other:?}"),
630        }
631    }
632}
lash_sansio/session_model/mod.rs

lash_sansio/session_model/
mod.rs
