1//! The in-memory reference adapter: every port, held in a `Mutex` and nothing else. It is what the
2//! conformance suite is first proven against (the suite must pass on the simplest honest
3//! implementation before it is allowed to judge another), what the service's tests run on, and a
4//! readable statement of what each port means.
5//!
6//! The outbound ports (`Judge`, `Embedder`, `Voice`, `Thinker`) are scripted: a test says what they
7//! answer and reads how often they were asked.
8//!
9//! Handles are cheap clones of shared state, so "the process restarted" is a clone of the handle: the
10//! state survives, which is exactly what a restart of a durable adapter must look like from outside.
11
12use std::collections::BTreeMap;
13use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
14use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
15
16use ::log::{debug, trace};
17use whiskers_core::{Age, ChatState, Direction, Household, IconId, MemoryDoc, MemorySnapshot, TokenLimit, Verdict};
18use whiskers_ports::{
19    Allowance, Appended, Audio, CreditsError, DeviceCursor, DeviceName, Diagnostic, EmbedBatch, EmbedError, Embedder, EntryLine,
20    Hub, Journal, Judge, JudgeError, Lookup, LogCursor, ManualClock, Merged, ModelName, Page, PictureShelf, Put, ReserveError, Reply,
21    SafePictureId, SecretName, Secrets, Settlement, SpeakError, SpeechLine, StoreError, StoredLine, ThinkError, ThinkRequest, Thinker,
22    ThinkingLedger, ThinkingRoom, ThinkingUsage, TimedAudio, TimingError, AudioMs, LineTiming, Tokens, Voice, VoiceCredits, VoiceDay, VoiceReservation, VoiceSpend,
23    Clock, HasAllowance, HasChat, HasClock, HasEmbedder, HasHousehold, HasJournal, HasJudge, HasMemory, HasPictures, HasThinker, HasVoice, Millis,
24};
25
26fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
27    m.lock().unwrap_or_else(PoisonError::into_inner)
28}
29
30/// The memory hub.
31#[derive(Clone, Default)]
32pub struct MemMemoryHub(Arc<Mutex<MemoryDoc>>);
33
34impl Hub for MemMemoryHub {
35    type Document = MemorySnapshot;
36
37    async fn merge(&self, theirs: MemorySnapshot) -> Result<Merged<MemorySnapshot>, StoreError> {
38        let mut doc = lock(&self.0);
39        let changed = doc.merge(theirs);
40        Ok(Merged { document: doc.snapshot(), changed })
41    }
42
43    async fn current(&self) -> Result<MemorySnapshot, StoreError> {
44        Ok(lock(&self.0).snapshot())
45    }
46}
47
48/// The household hub.
49#[derive(Clone, Default)]
50pub struct MemHouseholdHub(Arc<Mutex<Household>>);
51
52impl Hub for MemHouseholdHub {
53    type Document = Household;
54
55    async fn merge(&self, theirs: Household) -> Result<Merged<Household>, StoreError> {
56        let mut doc = lock(&self.0);
57        let changed = doc.merge(&theirs);
58        Ok(Merged { document: doc.clone(), changed: usize::from(changed) })
59    }
60
61    async fn current(&self) -> Result<Household, StoreError> {
62        Ok(lock(&self.0).clone())
63    }
64}
65
66/// The chat hub. Its storage can be told to refuse writes, to see what a failed save does.
67#[derive(Clone, Default)]
68pub struct MemChatHub {
69    doc: Arc<Mutex<ChatState>>,
70    refuse_writes: Arc<AtomicBool>,
71}
72
73impl MemChatHub {
74    /// From now on (until told otherwise) a copy that would have to be saved is refused.
75    pub fn refuse_writes(&self, refuse: bool) {
76        self.refuse_writes.store(refuse, Ordering::SeqCst);
77    }
78}
79
80impl Hub for MemChatHub {
81    type Document = ChatState;
82
83    async fn merge(&self, theirs: ChatState) -> Result<Merged<ChatState>, StoreError> {
84        let mut doc = lock(&self.doc);
85        if !theirs.is_newer_than(&doc) {
86            return Ok(Merged { document: doc.clone(), changed: 0 });
87        }
88        if self.refuse_writes.load(Ordering::SeqCst) {
89            debug!("mem chat hub: refusing to keep v{}", theirs.version);
90            return Err(StoreError::Unavailable);
91        }
92        let adopted = doc.merge(theirs);
93        Ok(Merged { document: doc.clone(), changed: usize::from(adopted) })
94    }
95
96    async fn current(&self) -> Result<ChatState, StoreError> {
97        Ok(lock(&self.doc).clone())
98    }
99}
100
101/// The journal.
102#[derive(Clone, Default)]
103pub struct MemJournal(Arc<Mutex<Vec<(DeviceName, StoredLine)>>>);
104
105impl MemJournal {
106    fn held_in(log: &[(DeviceName, StoredLine)], device: &DeviceName) -> DeviceCursor {
107        DeviceCursor::new(log.iter().filter(|(d, _)| d == device).count() as u64)
108    }
109}
110
111impl Journal for MemJournal {
112    async fn held(&self, device: &DeviceName) -> Result<DeviceCursor, StoreError> {
113        Ok(Self::held_in(&lock(&self.0), device))
114    }
115
116    async fn append(&self, device: &DeviceName, at: DeviceCursor, lines: &[EntryLine]) -> Result<Appended, StoreError> {
117        let mut log = lock(&self.0);
118        let held = Self::held_in(&log, device);
119        if lines.is_empty() {
120            return Ok(Appended::Accepted { held });
121        }
122        if at != held {
123            debug!("mem journal: {} asked to append at {} but holds {}", device.as_str(), at.get(), held.get());
124            return Ok(Appended::Misaligned { held });
125        }
126        for l in lines {
127            log.push((device.clone(), StoredLine::of(device, l)));
128        }
129        Ok(Appended::Accepted { held: held.after(lines.len()) })
130    }
131
132    async fn pull(&self, since: LogCursor) -> Result<Page, StoreError> {
133        let all: Vec<StoredLine> = lock(&self.0).iter().map(|(_, l)| l.clone()).collect();
134        Ok(Page::of(&all, since))
135    }
136}
137
138/// The shelf of pictures.
139#[derive(Clone, Default)]
140pub struct MemPictures(Arc<Mutex<BTreeMap<String, Vec<u8>>>>);
141
142impl PictureShelf for MemPictures {
143    async fn has(&self, id: &SafePictureId) -> Result<bool, StoreError> {
144        Ok(lock(&self.0).contains_key(id.as_str()))
145    }
146
147    async fn put(&self, id: &SafePictureId, bytes: &[u8]) -> Result<Put, StoreError> {
148        let mut shelf = lock(&self.0);
149        if shelf.contains_key(id.as_str()) {
150            trace!("mem pictures: {} already kept", id.as_str());
151            return Ok(Put::AlreadyKept);
152        }
153        shelf.insert(id.as_str().to_owned(), bytes.to_vec());
154        Ok(Put::Stored)
155    }
156
157    async fn get(&self, id: &SafePictureId) -> Result<Option<Vec<u8>>, StoreError> {
158        Ok(lock(&self.0).get(id.as_str()).cloned())
159    }
160}
161
162#[derive(Default)]
163struct AllowanceState {
164    thinking: ThinkingLedger,
165    voice: VoiceDay,
166    last_failure: Option<SpeakError>,
167}
168
169/// The allowance.
170#[derive(Clone, Default)]
171pub struct MemAllowance(Arc<Mutex<AllowanceState>>);
172
173impl Allowance for MemAllowance {
174    async fn thinking_room(&self, now: Millis, limit: &TokenLimit) -> Result<ThinkingRoom, StoreError> {
175        Ok(lock(&self.0).thinking.room(now, limit))
176    }
177
178    async fn charge_thinking(&self, now: Millis, cost: Tokens) -> Result<(), StoreError> {
179        lock(&self.0).thinking.charge(now, cost);
180        Ok(())
181    }
182
183    async fn thinking_usage(&self, now: Millis, limit: &TokenLimit) -> Result<ThinkingUsage, StoreError> {
184        Ok(lock(&self.0).thinking.usage(now, limit))
185    }
186
187    async fn reserve_voice(&self, now: Millis, chars: u32, cap: u32) -> Result<VoiceReservation, ReserveError> {
188        lock(&self.0).voice.reserve(now, chars, cap)
189    }
190
191    async fn settle_voice(&self, reservation: VoiceReservation, how: Settlement) -> Result<(), StoreError> {
192        let mut state = lock(&self.0);
193        match how {
194            Settlement::Spoken => state.last_failure = None,
195            Settlement::NotSpoken(why) => {
196                state.voice.release(reservation);
197                state.last_failure = Some(why);
198            }
199        }
200        Ok(())
201    }
202
203    async fn voice_spend(&self, now: Millis) -> Result<VoiceSpend, StoreError> {
204        let state = lock(&self.0);
205        Ok(VoiceSpend { spent_today: state.voice.spent_on(whiskers_ports::Day::containing(now)), last_failure: state.last_failure.clone() })
206    }
207}
208
209/// Secrets, as raw text exactly as a store would hold it.
210#[derive(Clone, Default)]
211pub struct MemSecrets(Arc<Mutex<BTreeMap<String, String>>>);
212
213impl MemSecrets {
214    pub fn with(entries: &[(SecretName, &str)]) -> Self {
215        let s = Self::default();
216        for (n, v) in entries {
217            s.set(*n, v);
218        }
219        s
220    }
221
222    pub fn set(&self, name: SecretName, raw: &str) {
223        lock(&self.0).insert(name.operator_name().to_owned(), raw.to_owned());
224    }
225}
226
227impl Secrets for MemSecrets {
228    async fn get(&self, name: SecretName) -> Result<Lookup, StoreError> {
229        Ok(Lookup::of(lock(&self.0).get(name.operator_name()).map(String::as_str)))
230    }
231}
232
233/// A Jev that answers what it is told to, and counts how often it was asked.
234pub struct ScriptedJudge {
235    pub verdict: Mutex<Result<Verdict, JudgeError>>,
236    pub ranking: Mutex<Result<Vec<f32>, JudgeError>>,
237    /// What `choose_icon` answers (the options it was last given are in `icon_options`).
238    pub icon: Mutex<Result<Option<IconId>, JudgeError>>,
239    pub icon_options: Mutex<Vec<IconId>>,
240    pub asked: AtomicUsize,
241}
242
243impl Default for ScriptedJudge {
244    fn default() -> Self {
245        Self {
246            verdict: Mutex::new(Ok(Verdict::Allow)),
247            ranking: Mutex::new(Ok(Vec::new())),
248            icon: Mutex::new(Ok(None)),
249            icon_options: Mutex::new(Vec::new()),
250            asked: AtomicUsize::new(0),
251        }
252    }
253}
254
255impl Judge for ScriptedJudge {
256    async fn check(&self, _direction: Direction, _age: Age, _text: &str) -> Result<Verdict, JudgeError> {
257        self.asked.fetch_add(1, Ordering::SeqCst);
258        lock(&self.verdict).clone()
259    }
260
261    async fn rerank(&self, _query: &str, candidates: &[String]) -> Result<Vec<f32>, JudgeError> {
262        self.asked.fetch_add(1, Ordering::SeqCst);
263        match &*lock(&self.ranking) {
264            Ok(p) if p.is_empty() => Ok(vec![0.5; candidates.len()]),
265            other => other.clone(),
266        }
267    }
268
269    async fn choose_icon(&self, _memory: &str, options: &[IconId]) -> Result<Option<IconId>, JudgeError> {
270        self.asked.fetch_add(1, Ordering::SeqCst);
271        *lock(&self.icon_options) = options.to_vec();
272        lock(&self.icon).clone()
273    }
274}
275
276/// How a scripted embedder turns the texts of a batch into vectors, when a test wants it to answer by what it was
277/// asked rather than with one fixed answer.
278pub type EmbedWith = std::sync::Arc<dyn Fn(&[String]) -> Vec<Vec<f32>> + Send + Sync>;
279
280/// An embedding model that answers what it is told to.
281pub struct ScriptedEmbedder {
282    pub answer: Mutex<Result<Vec<Vec<f32>>, EmbedError>>,
283    /// If set, used instead of `answer`.
284    pub with: Mutex<Option<EmbedWith>>,
285    pub asked: AtomicUsize,
286}
287
288impl Default for ScriptedEmbedder {
289    fn default() -> Self {
290        Self {
291            answer: Mutex::new(Err(EmbedError::Unreachable(Diagnostic::new("embedding server: not scripted")))),
292            with: Mutex::new(None),
293            asked: AtomicUsize::new(0),
294        }
295    }
296}
297
298impl Embedder for ScriptedEmbedder {
299    async fn embed(&self, batch: &EmbedBatch) -> Result<Vec<Vec<f32>>, EmbedError> {
300        self.asked.fetch_add(1, Ordering::SeqCst);
301        if let Some(f) = lock(&self.with).clone() {
302            return Ok(f(batch.texts()));
303        }
304        lock(&self.answer).clone()
305    }
306}
307
308/// A voice that answers what it is told to.
309pub struct ScriptedVoice {
310    pub configured: Mutex<bool>,
311    pub speech: Mutex<Result<Vec<u8>, SpeakError>>,
312    pub credits: Mutex<Result<VoiceCredits, CreditsError>>,
313    pub asked: AtomicUsize,
314}
315
316impl Default for ScriptedVoice {
317    fn default() -> Self {
318        Self {
319            configured: Mutex::new(true),
320            speech: Mutex::new(Ok(vec![0xFF, 0xFB, 1, 2])),
321            credits: Mutex::new(Ok(VoiceCredits { tier: Some("free".into()), used: Some(10), limit: Some(10_000), resets_at_unix: Some(1_800_000_000) })),
322            asked: AtomicUsize::new(0),
323        }
324    }
325}
326
327impl Voice for ScriptedVoice {
328    async fn configured(&self) -> bool {
329        *lock(&self.configured)
330    }
331
332    async fn speak(&self, _line: &SpeechLine) -> Result<Audio, SpeakError> {
333        self.asked.fetch_add(1, Ordering::SeqCst);
334        if !*lock(&self.configured) {
335            return Err(SpeakError::NotConfigured);
336        }
337        lock(&self.speech).clone().map(Audio)
338    }
339
340    /// The scripted audio, and every character said for 100 ms in turn.
341    async fn speak_timed(&self, line: &SpeechLine) -> Result<TimedAudio, SpeakError> {
342        self.asked.fetch_add(1, Ordering::SeqCst);
343        if !*lock(&self.configured) {
344            return Err(SpeakError::NotConfigured);
345        }
346        let audio = Audio(lock(&self.speech).clone()?);
347        let spans = (0..line.chars()).map(|i| (AudioMs::new(i * 100), AudioMs::new(i * 100 + 100))).collect();
348        let unreadable = |e: TimingError| SpeakError::Unreadable(Diagnostic::new(e.to_string()));
349        let timing = LineTiming::new(line.as_str().to_owned(), spans).map_err(unreadable)?;
350        TimedAudio::new(audio, timing, line).map_err(unreadable)
351    }
352
353    async fn credits(&self) -> Result<VoiceCredits, CreditsError> {
354        lock(&self.credits).clone()
355    }
356}
357
358/// A model that answers what it is told to.
359pub struct ScriptedThinker {
360    pub model: ModelName,
361    pub answer: Mutex<Result<Reply, ThinkError>>,
362    pub asked: AtomicUsize,
363}
364
365impl Default for ScriptedThinker {
366    fn default() -> Self {
367        Self {
368            model: ModelName::new("m").expect("a model name"),
369            answer: Mutex::new(Ok(Reply { status: 200, body: r#"{"content":[],"usage":{"input_tokens":70,"output_tokens":30}}"#.into() })),
370            asked: AtomicUsize::new(0),
371        }
372    }
373}
374
375impl Thinker for ScriptedThinker {
376    fn serves(&self) -> &ModelName {
377        &self.model
378    }
379
380    async fn think(&self, _request: &ThinkRequest) -> Result<Reply, ThinkError> {
381        self.asked.fetch_add(1, Ordering::SeqCst);
382        lock(&self.answer).clone()
383    }
384}
385
386/// Every port, in memory. The fields are public so a test can script the outbound ports and read the
387/// state the service left behind.
388#[derive(Default)]
389pub struct MemBackend {
390    pub memory: MemMemoryHub,
391    pub household: MemHouseholdHub,
392    pub chat: MemChatHub,
393    pub journal: MemJournal,
394    pub pictures: MemPictures,
395    pub allowance: MemAllowance,
396    pub judge: ScriptedJudge,
397    pub embedder: ScriptedEmbedder,
398    pub icons: whiskers_ports::IconVectors,
399    pub voice: ScriptedVoice,
400    pub thinker: ScriptedThinker,
401    pub clock: ManualClock,
402}
403
404// Every capability, so a `MemBackend` is a complete `Backend`.
405impl HasMemory for MemBackend {
406    type Memory = MemMemoryHub;
407    fn memory(&self) -> &MemMemoryHub {
408        &self.memory
409    }
410}
411impl HasHousehold for MemBackend {
412    type Household = MemHouseholdHub;
413    fn household(&self) -> &MemHouseholdHub {
414        &self.household
415    }
416}
417impl HasChat for MemBackend {
418    type Chat = MemChatHub;
419    fn chat(&self) -> &MemChatHub {
420        &self.chat
421    }
422}
423impl HasJournal for MemBackend {
424    type Journal = MemJournal;
425    fn journal(&self) -> &MemJournal {
426        &self.journal
427    }
428}
429impl HasPictures for MemBackend {
430    type Pictures = MemPictures;
431    fn pictures(&self) -> &MemPictures {
432        &self.pictures
433    }
434}
435impl HasAllowance for MemBackend {
436    type Allowance = MemAllowance;
437    fn allowance(&self) -> &MemAllowance {
438        &self.allowance
439    }
440}
441impl HasJudge for MemBackend {
442    type Judge = ScriptedJudge;
443    fn judge(&self) -> &ScriptedJudge {
444        &self.judge
445    }
446}
447impl whiskers_ports::HasIcons for MemBackend {
448    fn icon_vectors(&self) -> &whiskers_ports::IconVectors {
449        &self.icons
450    }
451}
452impl HasEmbedder for MemBackend {
453    type Embedder = ScriptedEmbedder;
454    fn embedder(&self) -> &ScriptedEmbedder {
455        &self.embedder
456    }
457}
458impl HasVoice for MemBackend {
459    type Voice = ScriptedVoice;
460    fn voice(&self) -> &ScriptedVoice {
461        &self.voice
462    }
463}
464impl HasThinker for MemBackend {
465    type Thinker = ScriptedThinker;
466    fn thinker(&self) -> &ScriptedThinker {
467        &self.thinker
468    }
469}
470impl HasClock for MemBackend {
471    type Clock = ManualClock;
472    fn clock(&self) -> &ManualClock {
473        &self.clock
474    }
475}
476
477impl MemBackend {
478    /// A backend whose clock reads `ms`.
479    pub fn at(ms: u64) -> Self {
480        let b = Self::default();
481        b.clock.set(Millis::new(ms));
482        b
483    }
484
485    pub fn now(&self) -> Millis {
486        self.clock.now()
487    }
488}
489
490/// The harness of the in-memory adapter: a fixture for every capability, so the whole suite runs on it
491/// (`Suite::complete(MemFixtures)`). A restart is the identity: what is in memory stays in memory, which is
492/// the most a process that keeps nothing on disk can promise.
493#[derive(Clone, Copy, Default)]
494pub struct MemFixtures;
495
496macro_rules! mem_stateful {
497    ($fixture:ident, $port:ident) => {
498        impl crate::suite::fixture::$fixture for MemFixtures {
499            type Port = $port;
500            fn fresh(&self) -> $port {
501                $port::default()
502            }
503            fn restart(&self, port: $port) -> $port {
504                port
505            }
506        }
507    };
508}
509
510mem_stateful!(MemoryFixture, MemMemoryHub);
511mem_stateful!(HouseholdFixture, MemHouseholdHub);
512mem_stateful!(ChatFixture, MemChatHub);
513mem_stateful!(JournalFixture, MemJournal);
514mem_stateful!(PicturesFixture, MemPictures);
515mem_stateful!(AllowanceFixture, MemAllowance);
516
517impl crate::suite::fixture::ChatFaultFixture for MemFixtures {
518    type Port = MemChatHub;
519    fn fresh_refusing(&self) -> (MemChatHub, crate::suite::hubs::WriteSwitch) {
520        let hub = MemChatHub::default();
521        let switch = hub.clone();
522        (hub, Box::new(move |refuse| switch.refuse_writes(refuse)))
523    }
524    fn restart_refusing(&self, port: MemChatHub) -> MemChatHub {
525        port
526    }
527}
528
529impl crate::suite::fixture::JudgeFixture for MemFixtures {
530    type Port = ScriptedJudge;
531    fn without_credentials(&self) -> ScriptedJudge {
532        let judge = ScriptedJudge::default();
533        *lock(&judge.verdict) = Err(JudgeError::NotConfigured(SecretName::JevKey));
534        *lock(&judge.ranking) = Err(JudgeError::NotConfigured(SecretName::JevKey));
535        *lock(&judge.icon) = Err(JudgeError::NotConfigured(SecretName::JevKey));
536        judge
537    }
538}
539
540impl crate::suite::fixture::VoiceFixture for MemFixtures {
541    type Port = ScriptedVoice;
542    fn without_credentials(&self) -> ScriptedVoice {
543        let voice = ScriptedVoice::default();
544        *lock(&voice.configured) = false;
545        voice
546    }
547}
548
549impl crate::suite::fixture::ClockFixture for MemFixtures {
550    type Port = ManualClock;
551    fn clock(&self) -> ManualClock {
552        // A manual clock that reads a plausible time (the law asks for one after 2020).
553        ManualClock::at(1_800_000_000_000)
554    }
555}