theway_contract/
session.rs

1//! Engine-independent session persistence contracts.
2
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7use crate::extension::ExtensionDurableEntry;
8
9#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum SessionErrorCode {
12    NotFound,
13    AlreadyExists,
14    Corrupted,
15    StorageFailure,
16    Aborted,
17    Unknown,
18}
19
20#[derive(Clone, Debug, thiserror::Error)]
21#[error("{message}")]
22pub struct SessionError {
23    pub code: SessionErrorCode,
24    pub message: String,
25}
26
27impl SessionError {
28    pub fn new(code: SessionErrorCode, message: impl Into<String>) -> Self {
29        Self {
30            code,
31            message: message.into(),
32        }
33    }
34
35    pub fn corrupted(message: impl Into<String>) -> Self {
36        Self::new(SessionErrorCode::Corrupted, message)
37    }
38}
39
40#[derive(Clone, Debug, Serialize, Deserialize)]
41pub struct SessionMetadata {
42    pub id: String,
43    #[serde(rename = "createdAt")]
44    pub created_at: String,
45}
46
47/// Persisted session runtime values; credentials must not be added here.
48#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
49#[serde(rename_all = "camelCase")]
50pub struct SessionRuntimeContext {
51    #[serde(rename = "workDir")]
52    pub work_dir: String,
53    #[serde(default, skip_serializing_if = "Option::is_none")]
54    pub provider: Option<String>,
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    pub model: Option<String>,
57    #[serde(default, skip_serializing_if = "Option::is_none", rename = "baseUrl")]
58    pub base_url: Option<String>,
59    #[serde(default, skip_serializing_if = "Option::is_none")]
60    pub thinking: Option<bool>,
61}
62
63/// Persisted client identity and its non-secret runtime context.
64#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
65#[serde(rename_all = "camelCase")]
66pub struct SessionBinding {
67    #[serde(rename = "clientKey")]
68    pub client_key: String,
69    pub runtime: SessionRuntimeContext,
70}
71
72#[derive(Clone, Debug, Serialize, Deserialize)]
73pub struct JsonlSessionMetadata {
74    #[serde(flatten)]
75    pub base: SessionMetadata,
76    pub cwd: String,
77    pub path: String,
78    #[serde(
79        default,
80        skip_serializing_if = "Option::is_none",
81        rename = "parentSessionPath"
82    )]
83    pub parent_session_path: Option<String>,
84    #[serde(
85        default,
86        skip_serializing_if = "Option::is_none",
87        rename = "importedFrom"
88    )]
89    pub imported_from: Option<SessionImportOrigin>,
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub binding: Option<SessionBinding>,
92    #[serde(
93        default,
94        skip_serializing_if = "Option::is_none",
95        rename = "collapseNodeId"
96    )]
97    pub collapse_node_id: Option<String>,
98    #[serde(default, skip_serializing_if = "Option::is_none")]
99    pub collapsed: Option<bool>,
100}
101
102#[derive(Clone, Debug, Serialize, Deserialize)]
103pub struct SessionImportOrigin {
104    #[serde(rename = "sessionId")]
105    pub session_id: String,
106    pub cwd: String,
107    #[serde(rename = "exportedAt")]
108    pub exported_at: String,
109    #[serde(rename = "thewayVersion")]
110    pub theway_version: String,
111}
112
113/// One persisted session entry. `payload` is the canonical full tagged JSON
114/// object; the remaining fields are validated indexes used by backends.
115#[derive(Clone, Debug, PartialEq)]
116pub struct StoredSessionEntry {
117    pub id: String,
118    pub parent_id: Option<String>,
119    pub entry_type: String,
120    pub timestamp: String,
121    pub payload: Value,
122}
123
124impl StoredSessionEntry {
125    pub fn from_payload(payload: Value) -> Result<Self, SessionError> {
126        let object = payload
127            .as_object()
128            .ok_or_else(|| SessionError::corrupted("session entry must be a JSON object"))?;
129        let string = |field: &str| -> Result<String, SessionError> {
130            object
131                .get(field)
132                .and_then(Value::as_str)
133                .filter(|value| !value.is_empty())
134                .map(str::to_string)
135                .ok_or_else(|| {
136                    SessionError::corrupted(format!("session entry has invalid {field}"))
137                })
138        };
139        let parent_id = match object.get("parentId") {
140            None | Some(Value::Null) => None,
141            Some(Value::String(value)) => Some(value.clone()),
142            Some(_) => {
143                return Err(SessionError::corrupted(
144                    "session entry has invalid parentId",
145                ));
146            }
147        };
148        let entry = Self {
149            id: string("id")?,
150            parent_id,
151            entry_type: string("type")?,
152            timestamp: string("timestamp")?,
153            payload,
154        };
155        entry.validate_shape()?;
156        Ok(entry)
157    }
158
159    pub fn leaf(
160        id: String,
161        parent_id: Option<String>,
162        timestamp: String,
163        target_id: Option<String>,
164    ) -> Result<Self, SessionError> {
165        Self::from_payload(serde_json::json!({
166            "type": "leaf",
167            "id": id,
168            "parentId": parent_id,
169            "timestamp": timestamp,
170            "targetId": target_id,
171        }))
172    }
173
174    pub fn session_info(
175        id: String,
176        parent_id: Option<String>,
177        timestamp: String,
178        name: String,
179    ) -> Result<Self, SessionError> {
180        Self::from_payload(serde_json::json!({
181            "type": "session_info",
182            "id": id,
183            "parentId": parent_id,
184            "timestamp": timestamp,
185            "name": name,
186        }))
187    }
188
189    pub fn collapse(
190        id: String,
191        parent_id: Option<String>,
192        timestamp: String,
193        source_session_id: String,
194        child_session_id: String,
195        compact_text: String,
196        raw_text_ref: String,
197        collapsed_at: String,
198    ) -> Result<Self, SessionError> {
199        Self::from_payload(serde_json::json!({
200            "type": "collapse",
201            "id": id,
202            "parentId": parent_id,
203            "timestamp": timestamp,
204            "sourceSessionId": source_session_id,
205            "childSessionId": child_session_id,
206            "compactText": compact_text,
207            "rawTextRef": raw_text_ref,
208            "collapsedAt": collapsed_at,
209        }))
210    }
211
212    pub fn session_graph_state(
213        id: String,
214        parent_id: Option<String>,
215        timestamp: String,
216        data: Value,
217    ) -> Result<Self, SessionError> {
218        Self::from_payload(serde_json::json!({
219            "type": "custom",
220            "id": id,
221            "parentId": parent_id,
222            "timestamp": timestamp,
223            "customType": "session_graph_state",
224            "data": data,
225        }))
226    }
227
228    pub fn extension(
229        id: String,
230        parent_id: Option<String>,
231        timestamp: String,
232        extension: ExtensionDurableEntry,
233    ) -> Result<Self, SessionError> {
234        extension
235            .validate()
236            .map_err(|error| SessionError::corrupted(error.to_string()))?;
237        Self::from_payload(serde_json::json!({
238            "type": "extension",
239            "id": id,
240            "parentId": parent_id,
241            "timestamp": timestamp,
242            "extension": extension,
243        }))
244    }
245
246    pub fn leaf_target_id(&self) -> Option<Option<&str>> {
247        (self.entry_type == "leaf").then(|| self.payload.get("targetId").and_then(Value::as_str))
248    }
249
250    pub fn label_update(&self) -> Option<(&str, Option<&str>)> {
251        if self.entry_type != "label" {
252            return None;
253        }
254        let target = self.payload.get("targetId")?.as_str()?;
255        let label = self.payload.get("label").and_then(Value::as_str);
256        Some((target, label))
257    }
258
259    pub fn extension_payload(&self) -> Result<Option<ExtensionDurableEntry>, SessionError> {
260        if self.entry_type != "extension" {
261            return Ok(None);
262        }
263        let extension =
264            self.payload.get("extension").cloned().ok_or_else(|| {
265                SessionError::corrupted("extension session entry has no envelope")
266            })?;
267        let extension: ExtensionDurableEntry = serde_json::from_value(extension)
268            .map_err(|error| SessionError::corrupted(error.to_string()))?;
269        extension
270            .validate()
271            .map_err(|error| SessionError::corrupted(error.to_string()))?;
272        Ok(Some(extension))
273    }
274
275    fn validate_shape(&self) -> Result<(), SessionError> {
276        let object = self.payload.as_object().expect("validated object");
277        let require_string = |field: &str| {
278            object
279                .get(field)
280                .and_then(Value::as_str)
281                .map(|_| ())
282                .ok_or_else(|| {
283                    SessionError::corrupted(format!(
284                        "{} session entry has invalid {field}",
285                        self.entry_type
286                    ))
287                })
288        };
289        match self.entry_type.as_str() {
290            "message" => {
291                if !object.get("message").is_some_and(Value::is_object) {
292                    return Err(SessionError::corrupted(
293                        "message session entry has invalid message",
294                    ));
295                }
296            }
297            "thinking_level_change" => require_string("thinkingLevel")?,
298            "model_change" => {
299                require_string("provider")?;
300                require_string("modelId")?;
301            }
302            "compaction" => {
303                require_string("summary")?;
304                require_string("firstKeptEntryId")?;
305                if !object.get("tokensBefore").is_some_and(Value::is_u64) {
306                    return Err(SessionError::corrupted(
307                        "compaction session entry has invalid tokensBefore",
308                    ));
309                }
310            }
311            "branch_summary" => {
312                require_string("fromId")?;
313                require_string("summary")?;
314            }
315            "extension" => {
316                self.extension_payload()?;
317            }
318            "custom" => {
319                require_string("customType")?;
320                if object.get("customType").and_then(Value::as_str) == Some("session_graph_state") {
321                    let data = object.get("data").ok_or_else(|| {
322                        SessionError::corrupted(
323                            "session_graph_state custom entry has no data object",
324                        )
325                    })?;
326                    if !data.is_object() {
327                        return Err(SessionError::corrupted(
328                            "session_graph_state custom entry has invalid data",
329                        ));
330                    }
331                    if !data.get("dags").is_some_and(Value::is_array) {
332                        return Err(SessionError::corrupted(
333                            "session_graph_state custom entry has invalid dags",
334                        ));
335                    }
336                    if !data.get("subagents").is_some_and(Value::is_array) {
337                        return Err(SessionError::corrupted(
338                            "session_graph_state custom entry has invalid subagents",
339                        ));
340                    }
341                }
342            }
343            "custom_message" => {
344                require_string("customType")?;
345                if !object.contains_key("content")
346                    || !object.get("display").is_some_and(Value::is_boolean)
347                {
348                    return Err(SessionError::corrupted(
349                        "custom_message session entry has invalid content or display",
350                    ));
351                }
352            }
353            "collapse" => {
354                require_string("sourceSessionId")?;
355                require_string("childSessionId")?;
356                require_string("compactText")?;
357                require_string("rawTextRef")?;
358                require_string("collapsedAt")?;
359            }
360            "label" => require_string("targetId")?,
361            "session_info" => {
362                if object.get("name").is_some_and(|value| !value.is_string()) {
363                    return Err(SessionError::corrupted(
364                        "session_info session entry has invalid name",
365                    ));
366                }
367            }
368            "leaf" => {
369                if object
370                    .get("targetId")
371                    .is_some_and(|value| !value.is_null() && !value.is_string())
372                {
373                    return Err(SessionError::corrupted(
374                        "leaf session entry has invalid targetId",
375                    ));
376                }
377            }
378            other => {
379                return Err(SessionError::corrupted(format!(
380                    "unknown session entry type {other}"
381                )));
382            }
383        }
384        Ok(())
385    }
386}
387
388/// Validate append order and return the active leaf after replay.
389pub fn validate_session_entries(
390    entries: &[StoredSessionEntry],
391) -> Result<Option<String>, SessionError> {
392    let mut seen = std::collections::HashSet::new();
393    let mut active_leaf_id = None;
394    for entry in entries {
395        if !seen.insert(entry.id.clone()) {
396            return Err(SessionError::corrupted(
397                "session transcript contains duplicate entry id",
398            ));
399        }
400        if let Some(parent) = &entry.parent_id
401            && !seen.contains(parent)
402        {
403            return Err(SessionError::corrupted(
404                "session transcript contains dangling parent reference",
405            ));
406        }
407        active_leaf_id = match entry.leaf_target_id() {
408            Some(Some(target)) if !seen.contains(target) => {
409                return Err(SessionError::corrupted(
410                    "session transcript contains dangling leaf target",
411                ));
412            }
413            Some(target) => target.map(str::to_string),
414            None => Some(entry.id.clone()),
415        };
416    }
417    Ok(active_leaf_id)
418}
419
420#[async_trait]
421pub trait SessionReader: Send + Sync {
422    async fn get_metadata_json(&self) -> Result<Value, SessionError>;
423    async fn get_leaf_id(&self) -> Result<Option<String>, SessionError>;
424    async fn get_entry(&self, id: &str) -> Result<Option<StoredSessionEntry>, SessionError>;
425    async fn get_entries(&self) -> Result<Vec<StoredSessionEntry>, SessionError>;
426    async fn get_path_to_root(
427        &self,
428        leaf_id: Option<&str>,
429    ) -> Result<Vec<StoredSessionEntry>, SessionError>;
430    async fn find_entries(&self, entry_type: &str)
431    -> Result<Vec<StoredSessionEntry>, SessionError>;
432    async fn get_label(&self, id: &str) -> Result<Option<String>, SessionError>;
433
434    /// Return one extension's entries on the selected branch in root-to-leaf
435    /// replay order. `None` selects the store's current active leaf.
436    async fn get_extension_entries(
437        &self,
438        extension_id: &str,
439        leaf_id: Option<&str>,
440    ) -> Result<Vec<StoredSessionEntry>, SessionError> {
441        let selected_leaf = match leaf_id {
442            Some(id) => Some(id.to_string()),
443            None => self.get_leaf_id().await?,
444        };
445        let path = self.get_path_to_root(selected_leaf.as_deref()).await?;
446        let mut entries = Vec::new();
447        for stored in path {
448            let Some(extension) = stored.extension_payload()? else {
449                continue;
450            };
451            if extension.extension_id == extension_id {
452                entries.push(stored);
453            }
454        }
455        Ok(entries)
456    }
457}
458
459#[async_trait]
460pub trait SessionStore: SessionReader {
461    async fn set_leaf_id(&self, id: Option<String>) -> Result<(), SessionError>;
462    async fn create_entry_id(&self) -> Result<String, SessionError>;
463    /// Atomically append a sequence of entries in the provided order. Either
464    /// every entry becomes visible or none of them does.
465    async fn append_entries(&self, entries: Vec<StoredSessionEntry>) -> Result<(), SessionError>;
466
467    /// Persist or clear a non-secret client binding. Backends that cannot
468    /// support binding updates must fail closed instead of silently succeeding.
469    async fn set_binding(&self, _binding: Option<SessionBinding>) -> Result<(), SessionError> {
470        Err(SessionError::new(
471            SessionErrorCode::StorageFailure,
472            "session store does not support binding updates",
473        ))
474    }
475
476    /// Persist or clear the session-graph collapse node id. Backends that
477    /// cannot support collapse metadata must fail closed instead of silently
478    /// succeeding.
479    async fn set_collapse_node_id(&self, _id: Option<String>) -> Result<(), SessionError> {
480        Err(SessionError::new(
481            SessionErrorCode::StorageFailure,
482            "session store does not support collapse metadata",
483        ))
484    }
485
486    /// Persist or clear the collapsed flag. Backends that cannot support
487    /// collapse metadata must fail closed instead of silently succeeding.
488    async fn set_collapsed(&self, _collapsed: bool) -> Result<(), SessionError> {
489        Err(SessionError::new(
490            SessionErrorCode::StorageFailure,
491            "session store does not support collapse metadata",
492        ))
493    }
494
495    async fn append_entry(&self, entry: StoredSessionEntry) -> Result<(), SessionError> {
496        self.append_entries(vec![entry]).await
497    }
498}
theway_contract/session.rs

theway_contract/
session.rs
