From 0163767171ba22d81b6db950cba0fb3c9a953c66 Mon Sep 17 00:00:00 2001
From: Ruben Fischer <ruben.fischer@onyva.de>
Date: Thu, 19 Feb 2026 20:19:13 +0100
Subject: [PATCH] Smaller telegram Feature fixes

---
 src/services/telegram_service.py | 37 ++++++++++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/services/telegram_service.py b/src/services/telegram_service.py
index 7c1be9a..2c5f831 100644
--- a/src/services/telegram_service.py
+++ b/src/services/telegram_service.py
@@ -265,6 +265,19 @@ class TelegramService:
             reply_markup=self._post_type_keyboard(active_types)
         )
 
+    # ==================== TOKEN LIMIT CHECK ====================
+
+    async def _check_token_limit(self, user_id: str, db) -> tuple[bool, str]:
+        """Check company token limit for the user. Returns (can_proceed, error_msg)."""
+        try:
+            profile = await db.get_profile(UUID(user_id))
+            if profile and profile.company_id:
+                can_proceed, error_msg, _, _ = await db.check_company_token_limit(profile.company_id)
+                return can_proceed, error_msg
+        except Exception as e:
+            logger.warning(f"Token limit check failed (allowing): {e}")
+        return True, ""
+
     # ==================== FEEDBACK HANDLING ====================
 
     async def _handle_feedback(self, chat_id: str, user_id: str, feedback: str, conv: dict, db) -> None:
@@ -277,6 +290,12 @@ class TelegramService:
             await self._clear_conv(chat_id)
             return
 
+        # Check token limit before calling the LLM
+        can_proceed, limit_msg = await self._check_token_limit(user_id, db)
+        if not can_proceed:
+            await self.send_message(chat_id, f"⚠️ {limit_msg}")
+            return
+
         await self.send_message(chat_id, "⏳ Überarbeite deinen Post...")
 
         try:
@@ -287,9 +306,16 @@ class TelegramService:
                 suggestion=feedback
             )
 
-            # Update post in DB if we have a post_id
+            # Update post in DB: append new version, increment iteration counter
             if post_id:
-                await db.update_generated_post(UUID(post_id), {"post_content": improved})
+                existing = await db.get_generated_post(UUID(post_id))
+                if existing:
+                    new_versions = list(existing.writer_versions) + [improved]
+                    await db.update_generated_post(UUID(post_id), {
+                        "post_content": improved,
+                        "writer_versions": new_versions,
+                        "iterations": existing.iterations + 1,
+                    })
 
             # Update conversation state
             conv["post_content"] = improved
@@ -340,6 +366,13 @@ class TelegramService:
             await self._clear_conv(chat_id)
             return
 
+        # Check token limit before calling the LLM
+        can_proceed, limit_msg = await self._check_token_limit(user_id, db)
+        if not can_proceed:
+            await self.edit_message(chat_id, message_id, f"⚠️ {limit_msg}")
+            await self._clear_conv(chat_id)
+            return
+
         # Edit the post-type selection message to show progress
         await self.edit_message(chat_id, message_id, "⏳ Erstelle deinen Post... Das kann einen Moment dauern.")