feat: add Telegram topic (message_thread_id) support

- Parse message_thread_id from incoming messages
- Include thread_id in sendMessage API calls
- Pass thread_id through message queue system
- Replies now go to the correct topic in supergroups

1 files changed, 35 insertions, 19 deletions

diff --git a/Omni/Agent/Telegram/Messages.hs b/Omni/Agent/Telegram/Messages.hs
index dfa3a3d..eab9668 100644
--- a/Omni/Agent/Telegram/Messages.hs
+++ b/Omni/Agent/Telegram/Messages.hs
@@ -128,6 +128,7 @@ data ScheduledMessage = ScheduledMessage
   { smId :: Text,
     smUserId :: Maybe Text,
     smChatId :: Int,
+    smThreadId :: Maybe Int,
     smContent :: Text,
     smSendAt :: UTCTime,
     smCreatedAt :: UTCTime,
@@ -147,6 +148,7 @@ instance Aeson.ToJSON ScheduledMessage where
       [ "id" .= smId m,
         "user_id" .= smUserId m,
         "chat_id" .= smChatId m,
+        "thread_id" .= smThreadId m,
         "content" .= smContent m,
         "send_at" .= smSendAt m,
         "created_at" .= smCreatedAt m,
@@ -164,6 +166,7 @@ instance SQL.FromRow ScheduledMessage where
     id' <- SQL.field
     userId <- SQL.field
     chatId <- SQL.field
+    threadId <- SQL.field
     content <- SQL.field
     sendAt <- SQL.field
     createdAt <- SQL.field
@@ -180,6 +183,7 @@ instance SQL.FromRow ScheduledMessage where
         { smId = id',
           smUserId = userId,
           smChatId = chatId,
+          smThreadId = threadId,
           smContent = content,
           smSendAt = sendAt,
           smCreatedAt = createdAt,
@@ -199,13 +203,14 @@ maxRetries :: Int
 maxRetries = 5
 
 initScheduledMessagesTable :: SQL.Connection -> IO ()
-initScheduledMessagesTable conn =
+initScheduledMessagesTable conn = do
   SQL.execute_
     conn
     "CREATE TABLE IF NOT EXISTS scheduled_messages (\
     \  id TEXT PRIMARY KEY,\
     \  user_id TEXT,\
     \  chat_id INTEGER NOT NULL,\
+    \  thread_id INTEGER,\
     \  content TEXT NOT NULL,\
     \  send_at TIMESTAMP NOT NULL,\
     \  created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,\
@@ -217,16 +222,25 @@ initScheduledMessagesTable conn =
     \  correlation_id TEXT,\
     \  telegram_message_id INTEGER\
     \)"
+  migrateAddThreadId conn
+
+migrateAddThreadId :: SQL.Connection -> IO ()
+migrateAddThreadId conn = do
+  result <- try @SomeException <| SQL.execute_ conn "ALTER TABLE scheduled_messages ADD COLUMN thread_id INTEGER"
+  case result of
+    Left _ -> pure ()
+    Right () -> pure ()
 
 queueMessage ::
   Maybe Text ->
   Int ->
+  Maybe Int ->
   Text ->
   UTCTime ->
   Maybe Text ->
   Maybe Text ->
   IO Text
-queueMessage mUserId chatId content sendAt msgType correlationId = do
+queueMessage mUserId chatId mThreadId content sendAt msgType correlationId = do
   uuid <- UUID.nextRandom
   now <- getCurrentTime
   let msgId = UUID.toText uuid
@@ -235,34 +249,36 @@ queueMessage mUserId chatId content sendAt msgType correlationId = do
     SQL.execute
       conn
       "INSERT INTO scheduled_messages \
-      \(id, user_id, chat_id, content, send_at, created_at, status, retry_count, message_type, correlation_id) \
-      \VALUES (?, ?, ?, ?, ?, ?, 'pending', 0, ?, ?)"
-      (msgId, mUserId, chatId, content, sendAt, now, msgType, correlationId)
+      \(id, user_id, chat_id, thread_id, content, send_at, created_at, status, retry_count, message_type, correlation_id) \
+      \VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', 0, ?, ?)"
+      (msgId, mUserId, chatId, mThreadId, content, sendAt, now, msgType, correlationId)
   pure msgId
 
 enqueueImmediate ::
   Maybe Text ->
   Int ->
+  Maybe Int ->
   Text ->
   Maybe Text ->
   Maybe Text ->
   IO Text
-enqueueImmediate mUserId chatId content msgType correlationId = do
+enqueueImmediate mUserId chatId mThreadId content msgType correlationId = do
   now <- getCurrentTime
-  queueMessage mUserId chatId content now msgType correlationId
+  queueMessage mUserId chatId mThreadId content now msgType correlationId
 
 enqueueDelayed ::
   Maybe Text ->
   Int ->
+  Maybe Int ->
   Text ->
   NominalDiffTime ->
   Maybe Text ->
   Maybe Text ->
   IO Text
-enqueueDelayed mUserId chatId content delay msgType correlationId = do
+enqueueDelayed mUserId chatId mThreadId content delay msgType correlationId = do
   now <- getCurrentTime
   let sendAt = addUTCTime delay now
-  queueMessage mUserId chatId content sendAt msgType correlationId
+  queueMessage mUserId chatId mThreadId content sendAt msgType correlationId
 
 fetchDueMessages :: UTCTime -> Int -> IO [ScheduledMessage]
 fetchDueMessages now batchSize =
@@ -270,7 +286,7 @@ fetchDueMessages now batchSize =
     initScheduledMessagesTable conn
     SQL.query
       conn
-      "SELECT id, user_id, chat_id, content, send_at, created_at, status, \
+      "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \
       \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \
       \FROM scheduled_messages \
       \WHERE status = 'pending' AND send_at <= ? \
@@ -286,7 +302,7 @@ listPendingMessages mUserId chatId =
       Just uid ->
         SQL.query
           conn
-          "SELECT id, user_id, chat_id, content, send_at, created_at, status, \
+          "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \
           \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \
           \FROM scheduled_messages \
           \WHERE user_id = ? AND chat_id = ? AND status = 'pending' AND send_at > datetime('now') \
@@ -295,7 +311,7 @@ listPendingMessages mUserId chatId =
       Nothing ->
         SQL.query
           conn
-          "SELECT id, user_id, chat_id, content, send_at, created_at, status, \
+          "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \
           \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \
           \FROM scheduled_messages \
           \WHERE chat_id = ? AND status = 'pending' AND send_at > datetime('now') \
@@ -309,7 +325,7 @@ getMessageById msgId =
     results <-
       SQL.query
         conn
-        "SELECT id, user_id, chat_id, content, send_at, created_at, status, \
+        "SELECT id, user_id, chat_id, thread_id, content, send_at, created_at, status, \
         \retry_count, last_attempt_at, last_error, message_type, correlation_id, telegram_message_id \
         \FROM scheduled_messages \
         \WHERE id = ?"
@@ -380,7 +396,7 @@ cancelMessage msgId =
     changes <- SQL.changes conn
     pure (changes > 0)
 
-messageDispatchLoop :: (Int -> Text -> IO (Maybe Int)) -> IO ()
+messageDispatchLoop :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> IO ()
 messageDispatchLoop sendFn =
   forever <| do
     now <- getCurrentTime
@@ -391,11 +407,11 @@ messageDispatchLoop sendFn =
         forM_ due <| \m -> dispatchOne sendFn m
         when (length due < 10) <| threadDelay 1000000
 
-dispatchOne :: (Int -> Text -> IO (Maybe Int)) -> ScheduledMessage -> IO ()
+dispatchOne :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> ScheduledMessage -> IO ()
 dispatchOne sendFn m = do
   now <- getCurrentTime
   markSending (smId m) now
-  result <- try (sendFn (smChatId m) (smContent m))
+  result <- try (sendFn (smChatId m) (smThreadId m) (smContent m))
   case result of
     Left (e :: SomeException) -> do
       let err = "Exception sending Telegram message: " <> tshow e
@@ -409,8 +425,8 @@ dispatchOne sendFn m = do
       markSent (smId m) (Just telegramMsgId) now'
       putText <| "Sent message " <> smId m <> " -> telegram_id " <> tshow telegramMsgId
 
-sendMessageTool :: Text -> Int -> Engine.Tool
-sendMessageTool uid chatId =
+sendMessageTool :: Text -> Int -> Maybe Int -> Engine.Tool
+sendMessageTool uid chatId mThreadId =
   Engine.Tool
     { Engine.toolName = "send_message",
       Engine.toolDescription =
@@ -453,7 +469,7 @@ sendMessageTool uid chatId =
                 let delay = fromIntegral (fromMaybe 0 delaySeconds)
                 now <- getCurrentTime
                 let sendAt = addUTCTime delay now
-                msgId <- queueMessage (Just uid) chatId text sendAt (Just "agent_tool") Nothing
+                msgId <- queueMessage (Just uid) chatId mThreadId text sendAt (Just "agent_tool") Nothing
                 pure
                   <| Aeson.object
                     [ "status" .= ("queued" :: Text),