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

lash_sansio/session_model/
mod.rs
