Skip to main content

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, 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#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
97pub struct ErrorEnvelope {
98    pub kind: String,
99    #[serde(default, skip_serializing_if = "Option::is_none")]
100    pub code: Option<String>,
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub terminal_reason: Option<crate::llm::types::LlmTerminalReason>,
103    pub user_message: String,
104    #[serde(default, skip_serializing_if = "Option::is_none")]
105    pub raw: Option<String>,
106}
107
108#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
109#[serde(tag = "type")]
110#[allow(clippy::large_enum_variant)]
111pub enum SessionEvent {
112    #[serde(rename = "text_delta")]
113    TextDelta { content: String },
114    /// Streaming update for the model's reasoning summary ("thinking").
115    /// The UI renders these incrementally in a muted/italic style;
116    /// reasoning content is never fed back to the model on subsequent
117    /// turns.
118    #[serde(rename = "reasoning_delta")]
119    ReasoningDelta { content: String },
120    #[serde(rename = "tool_call")]
121    ToolCall {
122        #[serde(default, skip_serializing_if = "Option::is_none")]
123        call_id: Option<String>,
124        name: String,
125        args: serde_json::Value,
126        output: crate::ToolCallOutput,
127        duration_ms: u64,
128    },
129    #[serde(rename = "tool_call_start")]
130    ToolCallStart {
131        #[serde(default, skip_serializing_if = "Option::is_none")]
132        call_id: Option<String>,
133        name: String,
134        args: serde_json::Value,
135    },
136    #[serde(rename = "message")]
137    Message { text: String, kind: String },
138    #[serde(rename = "llm_request")]
139    LlmRequest {
140        protocol_iteration: usize,
141        message_count: usize,
142        tool_list: String,
143    },
144    #[serde(rename = "llm_response")]
145    LlmResponse {
146        protocol_iteration: usize,
147        content: String,
148        duration_ms: u64,
149    },
150    #[serde(rename = "token_usage")]
151    TokenUsage {
152        protocol_iteration: usize,
153        usage: TokenUsage,
154        cumulative: TokenUsage,
155    },
156    #[serde(rename = "child_token_usage")]
157    ChildTokenUsage {
158        session_id: String,
159        source: String,
160        model: String,
161        protocol_iteration: usize,
162        usage: TokenUsage,
163        cumulative: TokenUsage,
164    },
165    #[serde(rename = "retry_status")]
166    RetryStatus {
167        wait_seconds: u64,
168        attempt: usize,
169        max_attempts: usize,
170        reason: String,
171        #[serde(default, skip_serializing_if = "Option::is_none")]
172        envelope: Option<ErrorEnvelope>,
173    },
174    #[serde(rename = "injected_turn_input_accepted")]
175    InjectedTurnInputAccepted {
176        inputs: Vec<AcceptedInjectedTurnInput>,
177        checkpoint: CheckpointKind,
178    },
179    #[serde(rename = "injected_messages_committed")]
180    InjectedMessagesCommitted {
181        messages: Vec<PluginMessage>,
182        checkpoint: CheckpointKind,
183    },
184    #[serde(rename = "plugin_event")]
185    PluginEvent {
186        plugin_id: String,
187        event: PluginRuntimeEvent,
188    },
189    /// Semantic result for a completed turn. `Done` remains the machine
190    /// lifecycle marker emitted after this event.
191    #[serde(rename = "turn_outcome")]
192    TurnOutcome { outcome: TurnOutcome },
193    #[serde(rename = "done")]
194    Done,
195    #[serde(rename = "error")]
196    Error {
197        message: String,
198        #[serde(default, skip_serializing_if = "Option::is_none")]
199        envelope: Option<ErrorEnvelope>,
200    },
201}
202
203#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
204#[serde(rename_all = "snake_case")]
205pub enum TurnOutcome {
206    Finished(TurnFinish),
207    AgentFrameSwitch { frame_id: String, task: String },
208    Stopped(TurnStop),
209}
210
211#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
212#[serde(rename_all = "snake_case")]
213pub enum TurnFinish {
214    AssistantMessage {
215        text: String,
216    },
217    FinalValue {
218        value: serde_json::Value,
219    },
220    ToolValue {
221        tool_name: String,
222        value: serde_json::Value,
223    },
224}
225
226#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
227#[serde(rename_all = "snake_case")]
228pub enum TurnStop {
229    Cancelled,
230    Incomplete,
231    InvalidInput,
232    MaxTurns,
233    ToolFailure,
234    ProviderError,
235    PluginAbort,
236    RuntimeError,
237    SubmittedError {
238        value: serde_json::Value,
239    },
240    ToolError {
241        tool_name: String,
242        value: serde_json::Value,
243    },
244}
245
246#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
247pub struct TurnTerminationPolicyState {
248    turn_limit_final_scheduled: bool,
249}
250
251impl Default for TurnTerminationPolicyState {
252    fn default() -> Self {
253        Self::new()
254    }
255}
256
257impl TurnTerminationPolicyState {
258    pub fn new() -> Self {
259        Self {
260            turn_limit_final_scheduled: false,
261        }
262    }
263
264    pub fn should_force_exit_after_grace_turn(&self) -> bool {
265        self.turn_limit_final_scheduled
266    }
267
268    pub fn turn_limit_final_to_schedule(
269        &self,
270        protocol_iteration: usize,
271        protocol_run_offset: usize,
272        max_turns: Option<usize>,
273    ) -> Option<usize> {
274        if self.turn_limit_final_scheduled {
275            return None;
276        }
277        let max = max_turns?;
278        if protocol_iteration < protocol_run_offset + max {
279            return None;
280        }
281        Some(max)
282    }
283
284    pub fn mark_turn_limit_final_scheduled(&mut self) {
285        self.turn_limit_final_scheduled = true;
286    }
287}
288
289pub fn make_error_envelope(
290    kind: &str,
291    code: Option<&str>,
292    terminal_reason: Option<crate::llm::types::LlmTerminalReason>,
293    user_message: impl Into<String>,
294    raw: Option<String>,
295) -> ErrorEnvelope {
296    let user_message = user_message.into();
297    ErrorEnvelope {
298        kind: kind.to_string(),
299        code: code.map(str::to_string),
300        terminal_reason,
301        user_message,
302        raw: raw.map(|s| truncate_raw_error(s.trim())),
303    }
304}
305
306pub fn make_error_event(
307    kind: &str,
308    code: Option<&str>,
309    user_message: impl Into<String>,
310    raw: Option<String>,
311) -> SessionEvent {
312    let user_message = user_message.into();
313    SessionEvent::Error {
314        message: user_message.clone(),
315        envelope: Some(make_error_envelope(kind, code, None, user_message, raw)),
316    }
317}
318
319pub fn truncate_raw_error(s: &str) -> String {
320    const MAX_RAW: usize = 4000;
321    let raw_len = s.chars().count();
322    if raw_len <= MAX_RAW {
323        return s.to_string();
324    }
325    let keep = MAX_RAW / 2;
326    let head = s.chars().take(keep).collect::<String>();
327    let tail = s
328        .chars()
329        .rev()
330        .take(keep)
331        .collect::<Vec<_>>()
332        .into_iter()
333        .rev()
334        .collect::<String>();
335    let omitted = raw_len.saturating_sub(keep * 2);
336    format!("{head}\n\n... ({omitted} chars omitted) ...\n\n{tail}")
337}
338
339pub fn reassign_part_ids(message_id: &str, parts: &mut [Part]) {
340    for (idx, part) in parts.iter_mut().enumerate() {
341        part.id = format!("{message_id}.p{idx}");
342    }
343}
344
345pub fn model_tool_specs_iter<'a>(
346    tools: impl IntoIterator<Item = &'a ToolDefinition>,
347) -> Vec<LlmToolSpec> {
348    tools
349        .into_iter()
350        .map(|tool| {
351            let model_tool = tool.model_tool();
352            LlmToolSpec {
353                name: model_tool.name,
354                description: model_tool.description,
355                input_schema: model_tool.input_schema,
356                output_schema: model_tool.output_schema,
357            }
358        })
359        .collect()
360}
361
362pub fn model_tool_specs(tools: &[ToolDefinition]) -> Vec<LlmToolSpec> {
363    model_tool_specs_iter(tools.iter())
364}