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#[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 #[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 #[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}