1use std::time::Duration;
13
14use okbot_core::{
15 AgentRunRecord, ChannelId, Clock, ConversationKey, ConversationStore, MessageRole, RunId,
16 ToolPolicy, UserId,
17};
18
19use crate::{ManualClock, bot_message, conversation, user_message};
20
21pub async fn conversation_store(store: &dyn ConversationStore, clock: &ManualClock) {
26 let a = conversation("zalo:shop", "room-a");
27 let b = conversation("zalo:shop", "room-b");
28 dedupe(store, clock).await;
29 pairing(store, clock).await;
30 fires_and_direct_chats(store, clock).await;
31
32 for i in 0..5 {
34 store
35 .append(user_message(&a, &format!("q{i}"), clock.now()))
36 .await
37 .expect("append");
38 store
39 .append(bot_message(&a, &format!("r{i}"), clock.now()))
40 .await
41 .expect("append");
42 }
43 store
44 .append(user_message(&b, "other room", clock.now()))
45 .await
46 .expect("append");
47 let recent = store.recent(&a, 3).await.expect("recent");
48 let texts: Vec<&str> = recent.iter().map(|m| m.text.as_str()).collect();
49 assert_eq!(texts, ["r3", "q4", "r4"]);
50 assert_eq!(recent[2].role, MessageRole::Bot);
51 assert_eq!(store.recent(&b, 10).await.expect("recent").len(), 1);
52 assert_eq!(
53 store
54 .recent(&conversation("zalo:shop", "none"), 10)
55 .await
56 .expect("recent"),
57 []
58 );
59
60 assert_eq!(store.session(&a).await.expect("session"), None);
62 store
63 .set_session(&a, "claude", "s-1")
64 .await
65 .expect("set_session");
66 store
67 .set_session(&a, "claude", "s-2")
68 .await
69 .expect("set_session");
70 assert_eq!(
71 store.session(&a).await.expect("session"),
72 Some(("claude".into(), "s-2".into()))
73 );
74 assert_eq!(store.session(&b).await.expect("session"), None);
75
76 store.reset(&a).await.expect("reset");
78 assert_eq!(
79 store.session(&a).await.expect("session"),
80 None,
81 "reset forgets the session"
82 );
83 assert_eq!(store.recent(&a, 10).await.expect("recent"), []);
84 store
85 .append(user_message(&a, "after reset", clock.now()))
86 .await
87 .expect("append");
88 assert_eq!(store.recent(&a, 10).await.expect("recent").len(), 1);
89 assert_eq!(store.recent(&b, 10).await.expect("recent").len(), 1);
90
91 store
93 .record_run(AgentRunRecord {
94 run_id: RunId::new("run-1"),
95 conversation: a.clone(),
96 backend: "fake".into(),
97 model: None,
98 tools: ToolPolicy::Web,
99 usage: None,
100 duration: Duration::from_millis(1500),
101 error: None,
102 started_at: clock.now(),
103 })
104 .await
105 .expect("record_run");
106
107 clock.advance(Duration::from_secs(3600));
109 store
110 .append(user_message(&b, "new", clock.now()))
111 .await
112 .expect("append");
113 let cutoff = clock
114 .now()
115 .checked_sub(jiff::SignedDuration::from_secs(60))
116 .expect("in range");
117 let removed = store.purge_before(cutoff).await.expect("purge");
118 assert!(removed >= 1, "purge removed nothing");
119 let left: Vec<String> = store
120 .recent(&b, 10)
121 .await
122 .expect("recent")
123 .into_iter()
124 .map(|m| m.text)
125 .collect();
126 assert_eq!(left, ["new"]);
127}
128
129async fn dedupe(store: &dyn ConversationStore, clock: &ManualClock) {
131 let channel: okbot_core::ChannelId = "zalo:shop".parse().expect("id");
132 let ttl = Duration::from_secs(60);
133 assert!(
134 store
135 .mark_seen(&channel, "e1", ttl)
136 .await
137 .expect("mark_seen")
138 );
139 assert!(
140 !store
141 .mark_seen(&channel, "e1", ttl)
142 .await
143 .expect("mark_seen")
144 );
145 let other: okbot_core::ChannelId = "zalo:other".parse().expect("id");
146 assert!(
147 store.mark_seen(&other, "e1", ttl).await.expect("mark_seen"),
148 "dedupe is per channel"
149 );
150 clock.advance(Duration::from_secs(61));
151 assert!(
152 store
153 .mark_seen(&channel, "e1", ttl)
154 .await
155 .expect("mark_seen"),
156 "expired mark"
157 );
158}
159
160pub async fn channel(channel: &dyn okbot_core::Channel, to: &okbot_core::ReplyTarget) {
171 use okbot_core::{ChannelError, RenderedMessage};
172
173 let caps = channel.caps().clone();
174 assert!(caps.max_message_len > 0, "max_message_len must be positive");
175 assert_eq!(
176 &to.conversation.channel,
177 channel.id(),
178 "the target belongs to this channel"
179 );
180
181 let sent = channel
182 .send(
183 to,
184 &RenderedMessage {
185 text: "conformance: hello".into(),
186 },
187 )
188 .await
189 .expect("send");
190 assert_eq!(&sent.target, to, "send returns the target it sent to");
191 assert!(
192 !sent.message_id.as_str().is_empty(),
193 "send returns the platform message id"
194 );
195
196 let longest = RenderedMessage {
197 text: "x".repeat(caps.max_message_len),
198 };
199 channel
200 .send(to, &longest)
201 .await
202 .expect("a message of exactly max_message_len is accepted");
203
204 let edited = channel
205 .edit(
206 &sent,
207 &RenderedMessage {
208 text: "conformance: edited".into(),
209 },
210 )
211 .await;
212 match (caps.edit, edited) {
213 (true, Ok(())) | (false, Err(ChannelError::Unsupported(_))) => {}
214 (true, Err(e)) => panic!("caps.edit is set but edit failed: {e}"),
215 (false, other) => panic!("caps.edit is unset, so edit must be Unsupported: {other:?}"),
216 }
217
218 channel
220 .typing(to)
221 .await
222 .expect("typing succeeds (a no-op when unsupported)");
223}
224
225pub async fn channel_stops(
230 channel: std::sync::Arc<dyn okbot_core::Channel>,
231 sink: std::sync::Arc<dyn okbot_core::InboundSink>,
232 shutdown: okbot_core::CancellationToken,
233) {
234 let run = tokio::spawn({
235 let shutdown = shutdown.clone();
236 async move { channel.run(sink, shutdown).await }
237 });
238 tokio::time::sleep(Duration::from_millis(200)).await;
239 shutdown.cancel();
240 tokio::time::timeout(Duration::from_secs(10), run)
241 .await
242 .expect("run returns within 10 s of shutdown")
243 .expect("run does not panic")
244 .expect("run returns Ok on shutdown");
245}
246
247async fn fires_and_direct_chats(store: &dyn ConversationStore, clock: &ManualClock) {
248 use okbot_core::FireState;
249 let at = clock.now();
250 let first = store.claim_fire("tasks/x/t-1", at).await.expect("claim");
251 assert_eq!((first.state, first.attempts), (FireState::Sending, 0));
252 store
253 .set_fire("tasks/x/t-1", at, FireState::Sending, 2)
254 .await
255 .expect("set_fire");
256 let again = store.claim_fire("tasks/x/t-1", at).await.expect("claim");
257 assert_eq!(again.attempts, 2, "a claim keeps the existing record");
258 store
259 .set_fire("tasks/x/t-1", at, FireState::Sent, 3)
260 .await
261 .expect("set_fire");
262 let sent = store.claim_fire("tasks/x/t-1", at).await.expect("claim");
263 assert_eq!(sent.state, FireState::Sent, "a sent firing stays sent");
264 let other = store
265 .claim_fire("tasks/x/t-1", at + jiff::SignedDuration::from_hours(24))
266 .await
267 .expect("claim");
268 assert_eq!(
269 other.state,
270 FireState::Sending,
271 "each firing time is its own record"
272 );
273 let recent = store.recent_fires(10).await.expect("recent_fires");
274 assert_eq!(recent.len(), 2);
275
276 let channel: ChannelId = "discord:main".parse().expect("valid id");
277 let user = UserId::new("u1");
278 assert_eq!(
279 store
280 .direct_chat(&channel, &user)
281 .await
282 .expect("direct_chat"),
283 None
284 );
285 let dm = ConversationKey::room(channel.clone(), "dm-1");
286 store
287 .set_direct_chat(&channel, &user, &dm)
288 .await
289 .expect("set_direct_chat");
290 let newer = ConversationKey::room(channel.clone(), "dm-2");
291 store
292 .set_direct_chat(&channel, &user, &newer)
293 .await
294 .expect("set_direct_chat");
295 assert_eq!(
296 store
297 .direct_chat(&channel, &user)
298 .await
299 .expect("direct_chat"),
300 Some(newer)
301 );
302}
303
304async fn pairing(store: &dyn ConversationStore, clock: &ManualClock) {
305 let zalo: okbot_core::ChannelId = "zalo:shop".parse().expect("valid id");
306 let discord: okbot_core::ChannelId = "discord:shop".parse().expect("valid id");
307 let me = okbot_core::UserId::new("me");
308 assert_eq!(store.owners().await.expect("owners"), vec![]);
309
310 store
311 .create_pairing_code("111111", Duration::from_secs(600))
312 .await
313 .expect("create");
314 assert!(
315 store
316 .redeem_pairing_code("111111", &zalo, &me)
317 .await
318 .expect("redeem"),
319 "a fresh code pairs"
320 );
321 assert!(
322 !store
323 .redeem_pairing_code("111111", &discord, &me)
324 .await
325 .expect("redeem"),
326 "codes are single-use"
327 );
328 assert_eq!(
329 store.owners().await.expect("owners"),
330 vec![(zalo.clone(), me.clone())]
331 );
332
333 store
335 .create_pairing_code("222222", Duration::from_secs(60))
336 .await
337 .expect("create");
338 clock.advance(Duration::from_secs(61));
339 assert!(
340 !store
341 .redeem_pairing_code("222222", &discord, &me)
342 .await
343 .expect("redeem")
344 );
345
346 store
348 .create_pairing_code("333333", Duration::from_secs(600))
349 .await
350 .expect("create");
351 for _ in 0..okbot_core::PAIRING_MAX_FAILURES {
352 assert!(
353 !store
354 .redeem_pairing_code("000000", &discord, &me)
355 .await
356 .expect("redeem")
357 );
358 }
359 assert!(
360 !store
361 .redeem_pairing_code("333333", &discord, &me)
362 .await
363 .expect("redeem"),
364 "cancelled after too many failures"
365 );
366
367 assert!(store.remove_owner(&zalo, &me).await.expect("remove"));
368 assert!(!store.remove_owner(&zalo, &me).await.expect("remove"));
369 assert_eq!(store.owners().await.expect("owners"), vec![]);
370}