From 4328fb3d2abbec91658f06c79d28caccf870b78c Mon Sep 17 00:00:00 2001
From: Rajat Ahuja <rahuja445@gmail.com>
Date: Fri, 5 Sep 2025 14:33:42 -0400
Subject: [PATCH] fix: actually pass in token count from msg to queue item

---
 src/deriver/enqueue.py       |  3 ++-
 src/deriver/queue_manager.py | 20 +++++++++-----------
 src/routers/messages.py      |  2 ++
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py
index 00028cda..5974fc9e 100644
--- a/src/deriver/enqueue.py
+++ b/src/deriver/enqueue.py
@@ -102,7 +102,6 @@ async def handle_session(
                 message_seq_map=message_seq_map,
             )
         )
-
     return queue_records
 
 
@@ -166,6 +165,7 @@ def create_representation_record(
         "payload": processed_payload,
         "session_id": session_id,
         "task_type": "representation",
+        "token_count": message.get("token_count"),
     }
 
 
@@ -198,6 +198,7 @@ def create_summary_record(
         "payload": processed_payload,
         "session_id": session_id,
         "task_type": "summary",
+        "token_count": message.get("token_count"),
     }
 
 
diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py
index cf902b53..6f2f683c 100644
--- a/src/deriver/queue_manager.py
+++ b/src/deriver/queue_manager.py
@@ -237,7 +237,8 @@ class QueueManager:
             try:
                 while not self.shutdown_event.is_set():
                     candidate_messages = await self.get_message_batch(
-                        work_unit_key, limit=20
+                        work_unit_key,
+                        limit=10,  # hard limit of 10 messages per batch
                     )
                     if not candidate_messages:
                         logger.debug(f"No more messages for work unit {work_unit_key}")
@@ -256,16 +257,13 @@ class QueueManager:
                         max_tokens = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
 
                         for msg in candidate_messages:
-                            if msg.task_type == "representation":
-                                msg_tokens = msg.token_count or 0
-                                if (
-                                    not messages_to_process
-                                    or token_count + msg_tokens <= max_tokens
-                                ):
-                                    messages_to_process.append(msg)
-                                    token_count += msg_tokens
-                                else:
-                                    break
+                            msg_tokens = msg.token_count or 0
+                            if (
+                                not messages_to_process
+                                or token_count + msg_tokens <= max_tokens
+                            ):
+                                messages_to_process.append(msg)
+                                token_count += msg_tokens
                             else:
                                 break
 
diff --git a/src/routers/messages.py b/src/routers/messages.py
index 3c59dbb3..27d89a14 100644
--- a/src/routers/messages.py
+++ b/src/routers/messages.py
@@ -66,6 +66,7 @@ async def create_messages_for_session(
                 "content": message.content,
                 "peer_name": message.peer_name,
                 "created_at": message.created_at,
+                "token_count": message.token_count,
             }
             for message in created_messages
         ]
@@ -127,6 +128,7 @@ async def create_messages_with_file(
             "content": message.content,
             "peer_name": message.peer_name,
             "created_at": message.created_at,
+            "token_count": message.token_count,
         }
         for message in created_messages
     ]