starweaver-session 0.9.0

Durable session contracts for Starweaver
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

use chrono::Utc;
use starweaver_core::{RunId, SessionId};
use starweaver_stream::{AgentStreamRecord, ReplayCursor, ReplayScope};

use crate::{
    error::{SessionStoreError, SessionStoreResult},
    records::StreamCursorRef,
};

use super::{InMemorySessionStore, run_key, run_key_label, store_failed};

impl InMemorySessionStore {
    pub(super) fn append_stream_record_batch(
        &self,
        session_id: &SessionId,
        run_id: &RunId,
        records: Vec<AgentStreamRecord>,
    ) -> SessionStoreResult<()> {
        let mut inner = self.inner.lock().map_err(store_failed)?;
        let key = run_key(session_id, run_id);
        if !inner.runs.contains_key(&key) {
            return Err(SessionStoreError::NotFound(run_key_label(
                session_id, run_id,
            )));
        }
        let stream = inner.streams.entry(key.clone()).or_default();
        for (index, record) in records.iter().enumerate() {
            let existing = stream
                .iter()
                .find(|existing| existing.sequence == record.sequence)
                .or_else(|| {
                    records[..index]
                        .iter()
                        .find(|existing| existing.sequence == record.sequence)
                });
            if existing.is_some_and(|existing| existing != record) {
                return Err(SessionStoreError::Failed(format!(
                    "stream record conflict for session {} run {} at sequence {}",
                    session_id.as_str(),
                    run_id.as_str(),
                    record.sequence
                )));
            }
        }
        for record in records {
            if stream
                .iter()
                .all(|existing| existing.sequence != record.sequence)
            {
                stream.push(record);
            }
        }
        stream.sort_by_key(|record| record.sequence);
        let last_sequence = stream.last().map(|record| record.sequence);
        if let Some(run) = inner.runs.get_mut(&key) {
            if let Some(sequence) = last_sequence {
                let cursor = StreamCursorRef::new(ReplayCursor::raw_runtime(
                    ReplayScope::run(run_id.as_str()),
                    sequence,
                ));
                run.stream_cursors
                    .retain(|existing| !existing.same_stream(&cursor));
                run.stream_cursors.push(cursor);
            }
            run.updated_at = Utc::now();
        }
        Ok(())
    }

    pub(super) fn replay_stream_record_batch(
        &self,
        session_id: &SessionId,
        run_id: &RunId,
    ) -> SessionStoreResult<Vec<AgentStreamRecord>> {
        let inner = self.inner.lock().map_err(store_failed)?;
        Ok(inner
            .streams
            .get(&run_key(session_id, run_id))
            .cloned()
            .unwrap_or_default())
    }

    pub(super) fn save_stream_cursor_ref(
        &self,
        session_id: &SessionId,
        run_id: &RunId,
        cursor: StreamCursorRef,
    ) -> SessionStoreResult<()> {
        cursor
            .validate_for_run(run_id)
            .map_err(|error| SessionStoreError::Failed(error.to_string()))?;
        let mut inner = self.inner.lock().map_err(store_failed)?;
        let run_key = run_key(session_id, run_id);
        let updated_at = Utc::now();
        let existing_run = inner
            .runs
            .get(&run_key)
            .ok_or_else(|| SessionStoreError::NotFound(run_key_label(session_id, run_id)))?;
        for existing in existing_run.stream_cursors.iter().chain(
            inner
                .sessions
                .get(session_id)
                .into_iter()
                .flat_map(|session| &session.stream_cursors),
        ) {
            cursor
                .validate_progression(existing)
                .map_err(|error| SessionStoreError::Failed(error.to_string()))?;
        }
        let run = inner
            .runs
            .get_mut(&run_key)
            .ok_or_else(|| SessionStoreError::NotFound(run_key_label(session_id, run_id)))?;
        run.stream_cursors
            .retain(|existing| !existing.same_stream(&cursor));
        run.stream_cursors.push(cursor.clone());
        run.updated_at = updated_at;
        if let Some(session) = inner.sessions.get_mut(session_id) {
            session
                .stream_cursors
                .retain(|existing| !existing.same_stream(&cursor));
            session.stream_cursors.push(cursor);
            session.updated_at = updated_at;
        }
        Ok(())
    }
}