aboutsummaryrefslogtreecommitdiffstats
path: root/crates/asist-app/src/session.rs
blob: 8d273c7ab2732bfc85ee86a9c0140c126fe7baf5 (plain)
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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
//! State shared between the pipeline threads.
//!
//! Only two things really need sharing: which turn is current, and the flags
//! that allow cutting whatever is in progress. Everything else travels over
//! channels. Keeping this surface small is what keeps interruption easy to
//! follow: cutting means bumping the turn and raising two flags.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;

use asist_core::event::TurnId;
use asist_core::http::Cancel;

#[derive(Clone)]
pub struct Session {
    inner: Arc<Inner>,
}

struct Inner {
    /// Turn being served. Any work from an earlier turn that shows up on a
    /// channel is garbage and is thrown away.
    current: AtomicU64,
    /// The assistant is talking (or about to).
    speaking: AtomicBool,
    /// Cierre solicitado.
    stopping: AtomicBool,
    /// Cuts the model answer in progress.
    llm: Cancel,
    /// Cuts the synthesis in progress.
    tts: Cancel,
}

impl Default for Session {
    fn default() -> Self {
        Self::new()
    }
}

impl Session {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Inner {
                current: AtomicU64::new(0),
                speaking: AtomicBool::new(false),
                stopping: AtomicBool::new(false),
                llm: Cancel::new(),
                tts: Cancel::new(),
            }),
        }
    }

    pub fn current(&self) -> TurnId {
        TurnId(self.inner.current.load(Ordering::SeqCst))
    }

    /// Opens a new turn and returns its id.
    pub fn begin_turn(&self) -> TurnId {
        let id = TurnId(self.inner.current.fetch_add(1, Ordering::SeqCst) + 1);
        self.inner.llm.reset();
        self.inner.tts.reset();
        id
    }

    /// Is `turn` still the current turn?
    ///
    /// Threads check it before spending work: a sentence that belongs to an
    /// outdated turn must be neither synthesized nor heard.
    pub fn is_current(&self, turn: TurnId) -> bool {
        turn == self.current()
    }

    /// Cuts everything in progress for the current turn.
    pub fn interrupt(&self) {
        self.inner.llm.cancel();
        self.inner.tts.cancel();
        self.inner.speaking.store(false, Ordering::SeqCst);
    }

    pub fn llm_cancel(&self) -> &Cancel {
        &self.inner.llm
    }

    pub fn tts_cancel(&self) -> &Cancel {
        &self.inner.tts
    }

    pub fn set_speaking(&self, speaking: bool) {
        self.inner.speaking.store(speaking, Ordering::SeqCst);
    }

    pub fn is_speaking(&self) -> bool {
        self.inner.speaking.load(Ordering::SeqCst)
    }

    pub fn request_stop(&self) {
        self.inner.stopping.store(true, Ordering::SeqCst);
        self.interrupt();
    }

    pub fn is_stopping(&self) -> bool {
        self.inner.stopping.load(Ordering::SeqCst)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn each_turn_gets_an_increasing_id() {
        let session = Session::new();
        assert_eq!(session.begin_turn(), TurnId(1));
        assert_eq!(session.begin_turn(), TurnId(2));
        assert_eq!(session.current(), TurnId(2));
    }

    #[test]
    fn work_from_an_old_turn_becomes_stale() {
        let session = Session::new();
        let first = session.begin_turn();
        session.begin_turn();
        assert!(!session.is_current(first), "the old turn must be discarded");
    }

    #[test]
    fn opening_a_turn_rearms_cancellation() {
        let session = Session::new();
        session.begin_turn();
        session.interrupt();
        assert!(session.llm_cancel().is_cancelled());

        session.begin_turn();
        assert!(
            !session.llm_cancel().is_cancelled(),
            "a new turn must not inherit the previous cancellation"
        );
    }

    #[test]
    fn interrupting_silences_and_cancels_both_stages() {
        let session = Session::new();
        session.begin_turn();
        session.set_speaking(true);
        session.interrupt();
        assert!(session.tts_cancel().is_cancelled());
        assert!(!session.is_speaking());
    }

    #[test]
    fn requesting_shutdown_also_interrupts() {
        let session = Session::new();
        session.request_stop();
        assert!(session.is_stopping());
        assert!(session.llm_cancel().is_cancelled());
    }
}