1use std::any::Any;
2
3use botkit_core::action::AnyChatActionSender;
4use botkit_core::{ContextData, OptionValue};
5
6use crate::action::TelegramActionSender;
7use crate::client::TelegramClient;
8use crate::types::{EntityType, Update, UpdateKind};
9
10pub struct TelegramContextData {
12 pub update: Update,
13 client: TelegramClient,
15 channel_id: String,
18 chat_id: Option<i64>,
19 user_id: String,
20 user_name: String,
21 command_name: Option<String>,
22 command_args: Option<String>,
23}
24
25impl TelegramContextData {
26 pub fn new(update: Update, client: TelegramClient) -> Self {
27 let (chat_id, user, actor_chat) = match &update.kind {
28 UpdateKind::Message(m) | UpdateKind::EditedMessage(m) => {
29 (Some(m.chat.id), m.from.as_ref(), None)
30 }
31 UpdateKind::CallbackQuery(cq) => {
33 (cq.message.as_ref().map(|m| m.chat.id), Some(&cq.from), None)
34 }
35 UpdateKind::MessageReaction(r) => {
36 (Some(r.chat.id), r.user.as_ref(), r.actor_chat.as_ref())
37 }
38 UpdateKind::Unknown => (None, None, None),
39 };
40
41 let (user_id, user_name) = user
44 .map(|u| (u.id.to_string(), display_name(u)))
45 .or_else(|| {
46 actor_chat.map(|c| {
47 (
48 c.id.to_string(),
49 c.title
50 .clone()
51 .or_else(|| c.username.clone())
52 .unwrap_or_default(),
53 )
54 })
55 })
56 .unwrap_or_default();
57
58 let (command_name, command_args) = extract_command(&update);
59
60 Self {
61 channel_id: chat_id.map(|id| id.to_string()).unwrap_or_default(),
62 chat_id,
63 user_id,
64 user_name,
65 command_name,
66 command_args,
67 update,
68 client,
69 }
70 }
71
72 pub fn client(&self) -> &TelegramClient {
74 &self.client
75 }
76
77 pub fn chat_id(&self) -> Option<i64> {
79 self.chat_id
80 }
81
82 pub fn thread_id(&self) -> Option<i64> {
84 match &self.update.kind {
85 UpdateKind::Message(m) | UpdateKind::EditedMessage(m) => m.message_thread_id,
86 UpdateKind::CallbackQuery(cq) => cq.message.as_ref()?.message_thread_id,
87 UpdateKind::MessageReaction(_) | UpdateKind::Unknown => None,
88 }
89 }
90}
91
92fn display_name(user: &crate::types::User) -> String {
94 match &user.last_name {
95 Some(last) => format!("{} {last}", user.first_name),
96 None => user.first_name.clone(),
97 }
98}
99
100fn extract_command(update: &Update) -> (Option<String>, Option<String>) {
106 let UpdateKind::Message(message) = &update.kind else {
107 return (None, None);
108 };
109
110 let (Some(text), Some(entities)) = (&message.text, &message.entities) else {
111 return (None, None);
112 };
113
114 let entity = entities
116 .iter()
117 .find(|e| matches!(e.entity_type, EntityType::BotCommand) && e.offset == 0);
118
119 let Some(entity) = entity.filter(|e| e.length > 0) else {
120 return (None, None);
121 };
122
123 let split = usize::try_from(entity.length)
124 .ok()
125 .and_then(|length| utf16_offset_to_byte_index(text, length));
126
127 let Some(split) = split else {
128 return (None, None);
129 };
130
131 let (command, rest) = text.split_at(split);
132
133 let name = command
135 .trim_start_matches('/')
136 .split('@')
137 .next()
138 .unwrap_or_default();
139
140 if name.is_empty() {
141 return (None, None);
142 }
143
144 let args = rest.trim();
145 (
146 Some(name.to_string()),
147 (!args.is_empty()).then(|| args.to_string()),
148 )
149}
150
151fn utf16_offset_to_byte_index(text: &str, offset: usize) -> Option<usize> {
154 if offset == 0 {
155 return Some(0);
156 }
157
158 let mut units = 0;
159 for (index, ch) in text.char_indices() {
160 if units == offset {
161 return Some(index);
162 }
163 units += ch.len_utf16();
164 }
165
166 (units == offset).then_some(text.len())
167}
168
169impl ContextData for TelegramContextData {
170 fn channel_id(&self) -> &str {
171 &self.channel_id
172 }
173
174 fn user_id(&self) -> &str {
175 &self.user_id
176 }
177
178 fn user_name(&self) -> &str {
179 &self.user_name
180 }
181
182 fn command_name(&self) -> Option<&str> {
183 self.command_name.as_deref()
184 }
185
186 fn command_args(&self) -> Option<&str> {
187 self.command_args.as_deref()
188 }
189
190 fn option(&self, _name: &str) -> Option<OptionValue> {
191 None
193 }
194
195 fn button_id(&self) -> Option<&str> {
196 match &self.update.kind {
197 UpdateKind::CallbackQuery(cq) => cq.data.as_deref(),
198 _ => None,
199 }
200 }
201
202 fn message_content(&self) -> Option<&str> {
203 match &self.update.kind {
204 UpdateKind::Message(m) | UpdateKind::EditedMessage(m) => {
205 m.text.as_deref().or(m.caption.as_deref())
206 }
207 UpdateKind::CallbackQuery(cq) => cq.message.as_ref()?.text.as_deref().or(cq
208 .message
209 .as_ref()?
210 .caption
211 .as_deref()),
212 UpdateKind::MessageReaction(_) | UpdateKind::Unknown => None,
213 }
214 }
215
216 fn as_any(&self) -> &dyn Any {
217 self
218 }
219
220 fn action_sender(&self) -> Option<AnyChatActionSender> {
221 Some(AnyChatActionSender::new(TelegramActionSender::new(
222 self.client.clone(),
223 self.chat_id?,
224 self.thread_id(),
225 )))
226 }
227}
228
229#[cfg(test)]
230mod tests {
231 use super::*;
232
233 fn update(json: serde_json::Value) -> Update {
234 serde_json::from_value(json).expect("valid update")
235 }
236
237 fn message(text: &str, entities: serde_json::Value) -> Update {
238 update(serde_json::json!({
239 "update_id": 1,
240 "message": {
241 "message_id": 1,
242 "date": 0,
243 "chat": { "id": 42, "type": "private" },
244 "from": { "id": 7, "is_bot": false, "first_name": "Ada", "last_name": "Lovelace" },
245 "text": text,
246 "entities": entities,
247 }
248 }))
249 }
250
251 fn command_entity(length: i64) -> serde_json::Value {
252 serde_json::json!([{ "type": "bot_command", "offset": 0, "length": length }])
253 }
254
255 #[test]
256 fn parses_a_command_with_arguments() {
257 let update = message("/greet world and beyond", command_entity(6));
258 let (name, args) = extract_command(&update);
259 assert_eq!(name.as_deref(), Some("greet"));
260 assert_eq!(args.as_deref(), Some("world and beyond"));
261 }
262
263 #[test]
264 fn parses_a_bare_command() {
265 let (name, args) = extract_command(&message("/ping", command_entity(5)));
266 assert_eq!(name.as_deref(), Some("ping"));
267 assert_eq!(args, None);
268 }
269
270 #[test]
271 fn strips_the_bot_mention_suffix() {
272 let (name, args) = extract_command(&message("/ping@my_bot now", command_entity(12)));
273 assert_eq!(name.as_deref(), Some("ping"));
274 assert_eq!(args.as_deref(), Some("now"));
275 }
276
277 #[test]
278 fn ignores_commands_that_do_not_open_the_message() {
279 let update = message(
280 "see /ping",
281 serde_json::json!([{ "type": "bot_command", "offset": 4, "length": 5 }]),
282 );
283 assert_eq!(extract_command(&update), (None, None));
284 }
285
286 #[test]
287 fn ignores_messages_without_a_command_entity() {
288 let update = message("just chatting", serde_json::json!([]));
289 assert_eq!(extract_command(&update), (None, None));
290 }
291
292 #[test]
293 fn entity_lengths_are_counted_in_utf16_code_units() {
294 let update = message("/wave🎉 hi", command_entity(7));
297 let (name, args) = extract_command(&update);
298 assert_eq!(name.as_deref(), Some("wave🎉"));
299 assert_eq!(args.as_deref(), Some("hi"));
300 }
301
302 #[test]
303 fn out_of_range_entity_lengths_do_not_panic() {
304 for length in [-1, 0, 6, 9999] {
305 let update = message("/ping", command_entity(length));
306 assert_eq!(extract_command(&update), (None, None), "length {length}");
307 }
308 }
309
310 #[test]
311 fn entity_lengths_landing_mid_character_do_not_panic() {
312 let update = message("🎉x", command_entity(1));
314 assert_eq!(extract_command(&update), (None, None));
315 }
316
317 #[test]
318 fn callback_queries_expose_their_button_and_chat() {
319 let update = update(serde_json::json!({
320 "update_id": 2,
321 "callback_query": {
322 "id": "cb1",
323 "from": { "id": 7, "is_bot": false, "first_name": "Ada" },
324 "chat_instance": "x",
325 "data": "confirm_yes",
326 "message": {
327 "message_id": 1,
328 "date": 0,
329 "chat": { "id": 42, "type": "private" },
330 "text": "Are you sure?"
331 }
332 }
333 }));
334
335 let data = TelegramContextData::new(update, TelegramClient::new("token"));
336 assert_eq!(data.chat_id(), Some(42));
337 assert_eq!(data.button_id(), Some("confirm_yes"));
338 assert_eq!(data.command_name(), None);
339 assert_eq!(data.message_content(), Some("Are you sure?"));
340 assert_eq!(data.user_name(), "Ada");
341 }
342
343 #[test]
344 fn inline_callback_queries_have_no_chat_to_reply_in() {
345 let update = update(serde_json::json!({
346 "update_id": 3,
347 "callback_query": {
348 "id": "cb2",
349 "from": { "id": 7, "is_bot": false, "first_name": "Ada" },
350 "chat_instance": "x",
351 "inline_message_id": "inline-1",
352 "data": "x"
353 }
354 }));
355
356 let data = TelegramContextData::new(update, TelegramClient::new("token"));
357 assert_eq!(data.chat_id(), None);
358 assert_eq!(data.channel_id(), "");
359 assert!(data.action_sender().is_none());
360 }
361
362 #[test]
363 fn messages_expose_the_sender_and_chat() {
364 let data = TelegramContextData::new(
365 message("/greet you", command_entity(6)),
366 TelegramClient::new("token"),
367 );
368 assert_eq!(data.chat_id(), Some(42));
369 assert_eq!(data.channel_id(), "42");
370 assert_eq!(data.user_id(), "7");
371 assert_eq!(data.user_name(), "Ada Lovelace");
372 assert_eq!(data.command_name(), Some("greet"));
373 assert_eq!(data.command_args(), Some("you"));
374 assert!(data.action_sender().is_some());
375 }
376
377 #[test]
378 fn media_caption_serves_as_message_content() {
379 let update = update(serde_json::json!({
381 "update_id": 5,
382 "message": {
383 "message_id": 9,
384 "date": 0,
385 "chat": { "id": 42, "type": "private" },
386 "from": { "id": 7, "is_bot": false, "first_name": "Ada" },
387 "caption": "look at this",
388 "photo": [
389 { "file_id": "p1", "file_unique_id": "u1", "width": 90, "height": 90 },
390 { "file_id": "p2", "file_unique_id": "u2", "width": 800, "height": 600 }
391 ]
392 }
393 }));
394
395 let data = TelegramContextData::new(update, TelegramClient::new("token"));
396 assert_eq!(data.message_content(), Some("look at this"));
397 assert_eq!(data.chat_id(), Some(42));
398 }
399
400 #[test]
401 fn unsupported_update_kinds_are_inert() {
402 let update = update(serde_json::json!({
403 "update_id": 4,
404 "poll": { "id": "p1" }
405 }));
406
407 let data = TelegramContextData::new(update, TelegramClient::new("token"));
408 assert_eq!(data.chat_id(), None);
409 assert_eq!(data.user_id(), "");
410 assert_eq!(data.command_name(), None);
411 assert_eq!(data.button_id(), None);
412 assert_eq!(data.message_content(), None);
413 }
414}