feat: message ingestion forced-batching (#327)

* feat: implement forced batching for representation tasks and adjust max tokens

- Updated `REPRESENTATION_BATCH_MAX_TOKENS` to 1024 in `config.py`.
- Enhanced `get_and_claim_work_units` in `QueueManager` to enforce batching based on token thresholds.
- Added tests to ensure representation work units are only claimed when token counts meet or exceed the threshold.
- Introduced a synthesis prompt for tool execution to improve final response generation.

* chore: update queue-status docs to match new behavior

* fix: make representation work query efficient

* fix: Align alembic models and force dreams on for tests

* fix: clean queue between tests

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>

docs/v2.6.0-alpha/documentation/features/advanced/queue-status.mdx

@@ -4,11 +4,13 @@ description: Learn how to check the status of Honcho's reasoning
 icon: "lines-leaning"
 ---
 
-Whenever messages are stored in Honcho, a background process kicks off to [reason](/v2.6.0-alpha/documentation/core-concepts/reasoning) about the conversation and generate insights.
+Whenever messages are stored in Honcho, background processes kick off to [reason](/v2.6.0-alpha/documentation/core-concepts/reasoning) about the conversation and generate insights.
 
-Reasoning is an asynchronous process and, depending on load, may not immediately
-generate insights for the latest message you've sent. To help with this, Honcho
-provides several utilities to check the status of the queue.
+Reasoning is an asynchronous process and will not immediately
+generate insights for the latest message you've sent. This is
+by design: we want to reason efficiently over batches of messages
+rather than assessing each message in a vacuum. Honcho provides
+several utilities to check the status of the queue.
 
 <CodeGroup>
 ```python Python
@@ -75,7 +77,7 @@ work_units will be processed in parallel
   generate an additional work unit for every peer that has `observe_others=True`
 
 The `get_queue_status` and `poll_queue_status` methods can take additional
-parameters to scope the status to a specific work unit
+parameters to scope the status to a specific work unit:
 
 <CodeGroup>
 ```python Python
@@ -104,7 +106,7 @@ export const QueueStatusOptionsSchema = z.object({
 Additionally, there are queue status and polling queue status methods
 available on the session objects in each of the SDKs.
 
-Below are the function signatures for the session level queue status method
+Below are the function signatures for the session level queue status method:
 
 <CodeGroup>
 ```python python

src/config.py

@@ -267,8 +267,8 @@ class DeriverSettings(BackupLLMSettingsMixin, HonchoSettings):
 
     REPRESENTATION_BATCH_MAX_TOKENS: Annotated[
         int,
-        Field(default=4096, ge=128, le=16_384),
-    ] = 4096
+        Field(default=1024, ge=128, le=16_384),
+    ] = 1024
 
     @model_validator(mode="after")
     def validate_batch_tokens_vs_context_limit(self):
@@ -361,6 +361,7 @@ class DialecticSettings(HonchoSettings):
                 MODEL="gemini-3-flash-preview",
                 THINKING_BUDGET_TOKENS=0,
                 MAX_TOOL_ITERATIONS=5,
+                TOOL_CHOICE="any",
             ),
             "medium": DialecticLevelSettings(
                 PROVIDER="anthropic",

src/deriver/queue_manager.py

@@ -10,7 +10,7 @@ import sentry_sdk
 from dotenv import load_dotenv
 from nanoid import generate as generate_nanoid
 from sentry_sdk.integrations.asyncio import AsyncioIntegration
-from sqlalchemy import and_, delete, select, update
+from sqlalchemy import and_, delete, or_, select, update
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.sql import func
@@ -204,26 +204,64 @@ class QueueManager:
     async def get_and_claim_work_units(self) -> dict[str, str]:
         """
         Get available work units that aren't being processed.
+        For representation tasks, only returns work units with accumulated tokens
+        >= REPRESENTATION_BATCH_MAX_TOKENS (forced batching).
         Returns a dict mapping work_unit_key to aqs_id.
         """
         limit: int = max(0, self.workers - self.get_total_owned_work_units())
         if limit == 0:
             return {}
-        async with tracked_db(
-            "get_available_work_units"
-        ) as db:  # Get number of available workers
-            query = (
-                select(models.QueueItem.work_unit_key)
-                .limit(limit)
-                .outerjoin(
-                    models.ActiveQueueSession,
-                    models.QueueItem.work_unit_key
-                    == models.ActiveQueueSession.work_unit_key,
+
+        batch_max_tokens = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
+
+        async with tracked_db("get_available_work_units") as db:
+            representation_prefix = "representation:"
+            token_stats_subq = (
+                select(
+                    models.QueueItem.work_unit_key,
+                    func.sum(models.Message.token_count).label("total_tokens"),
+                )
+                .join(
+                    models.Message,
+                    models.QueueItem.message_id == models.Message.id,
                 )
                 .where(~models.QueueItem.processed)
-                .where(models.QueueItem.work_unit_key.isnot(None))
-                .where(models.ActiveQueueSession.work_unit_key.is_(None))
-                .distinct()
+                .where(models.QueueItem.work_unit_key.startswith(representation_prefix))
+                .group_by(models.QueueItem.work_unit_key)
+                .subquery()
+            )
+
+            work_units_subq = (
+                select(models.QueueItem.work_unit_key)
+                .where(~models.QueueItem.processed)
+                .group_by(models.QueueItem.work_unit_key)
+                .subquery()
+            )
+
+            query = (
+                select(work_units_subq.c.work_unit_key)
+                .limit(limit)
+                .outerjoin(
+                    token_stats_subq,
+                    work_units_subq.c.work_unit_key == token_stats_subq.c.work_unit_key,
+                )
+                .where(
+                    ~select(models.ActiveQueueSession.id)
+                    .where(
+                        models.ActiveQueueSession.work_unit_key
+                        == work_units_subq.c.work_unit_key
+                    )
+                    .exists()
+                )
+                .where(
+                    or_(
+                        ~work_units_subq.c.work_unit_key.startswith(
+                            representation_prefix
+                        ),
+                        func.coalesce(token_stats_subq.c.total_tokens, 0)
+                        >= batch_max_tokens,
+                    )
+                )
             )
 
             result = await db.execute(query)

src/models.py

@@ -473,6 +473,12 @@ class QueueItem(Base):
             "processed",
             "id",
         ),
+        Index(
+            "ux_queue_dream_pending_work_unit_key",
+            "work_unit_key",
+            unique=True,
+            postgresql_where=text("task_type = 'dream' AND processed = false"),
+        ),
     )
 
     def __repr__(self) -> str:

src/utils/clients.py

@@ -851,6 +851,16 @@ async def _execute_tool_loop(
         f"Tool execution loop reached max iterations ({max_tool_iterations})"
     )
 
+    # Add a synthesis prompt to help the model generate a response
+    # without tool calls - the conversation currently ends with tool results
+    # and the model may not know to produce text output
+    synthesis_prompt = (
+        "You have reached the maximum number of tool calls. "
+        "Based on all the information you have gathered, provide your final response now. "
+        "Do not attempt to call any more tools."
+    )
+    conversation_messages.append({"role": "user", "content": synthesis_prompt})
+
     # If streaming the final response, use the streaming helper with metadata
     if stream_final:
         stream = _stream_final_response(

tests/deriver/conftest.py

@@ -1,18 +1,35 @@
 import asyncio
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence
 from datetime import datetime, timezone
 from typing import Any, Literal, TypeAlias, cast
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from nanoid import generate as generate_nanoid
-from sqlalchemy import select
+from sqlalchemy import delete, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src import crud, models, schemas
 from src.utils.queue_payload import create_payload
 from src.utils.work_unit import construct_work_unit_key
 
+
+@pytest.fixture(autouse=True)
+async def clean_queue_tables(db_session: AsyncSession) -> AsyncGenerator[None, None]:
+    """Clean up queue-related tables before each test to ensure isolation.
+
+    This prevents webhook queue items and active queue sessions from previous tests
+    from polluting subsequent tests, which can cause issues when tests have a limit
+    on how many work units can be claimed (e.g., WORKERS=1).
+    """
+    # Clean up before the test
+    await db_session.execute(delete(models.ActiveQueueSession))
+    await db_session.execute(delete(models.QueueItem))
+    await db_session.commit()
+
+    yield
+
+
 QueuePayload: TypeAlias = dict[str, Any]
 QueuePayloadEntry: TypeAlias = QueuePayload | tuple[QueuePayload, int | None]
 

tests/deriver/test_queue_processing.py

@@ -1061,3 +1061,286 @@ class TestQueueProcessing:
             processed_batches[1]["payload_count"] == 1
         )  # Third message (exceeds limit)
         assert all(b["task_type"] == "representation" for b in processed_batches)
+
+    @pytest.mark.asyncio
+    async def test_forced_batching_waits_for_threshold(
+        self,
+        db_session: AsyncSession,
+        sample_session_with_peers: tuple[models.Session, list[models.Peer]],
+        create_queue_payload: Callable[..., Any],
+    ) -> None:
+        """Test that representation work units below token threshold are not claimed"""
+
+        session, peers = sample_session_with_peers
+        peer = peers[0]
+
+        # Create messages with tokens BELOW the threshold
+        limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
+        token_counts = [100, 100, 100]  # Total 300, way below 4096
+
+        messages: list[models.Message] = []
+        for i, token_count in enumerate(token_counts):
+            message = models.Message(
+                session_name=session.name,
+                workspace_name=session.workspace_name,
+                peer_name=peer.name,
+                content=f"Short message {i}",
+                token_count=token_count,
+                seq_in_session=i + 1,
+            )
+            db_session.add(message)
+            messages.append(message)
+
+        await db_session.commit()
+        for message in messages:
+            await db_session.refresh(message)
+
+        # Create queue items
+        queue_items: list[models.QueueItem] = []
+        for message in messages:
+            payload = create_queue_payload(
+                message=message,
+                task_type="representation",
+                observed=peer.name,
+                observer=peer.name,
+            )
+            work_unit_key = construct_work_unit_key(session.workspace_name, payload)
+
+            queue_item = models.QueueItem(
+                session_id=session.id,
+                task_type="representation",
+                work_unit_key=work_unit_key,
+                payload=payload,
+                processed=False,
+                workspace_name=session.workspace_name,
+                message_id=message.id,
+            )
+            db_session.add(queue_item)
+            queue_items.append(queue_item)
+
+        await db_session.commit()
+
+        qm = QueueManager()
+
+        # Try to claim work units - should NOT return the representation work unit
+        claimed = await qm.get_and_claim_work_units()
+
+        # The representation work unit should NOT be claimed (tokens below threshold)
+        rep_work_unit_key = queue_items[0].work_unit_key
+        assert rep_work_unit_key not in claimed
+
+        # Now add more messages to exceed the threshold
+        more_messages: list[models.Message] = []
+        for i in range(10):
+            message = models.Message(
+                session_name=session.name,
+                workspace_name=session.workspace_name,
+                peer_name=peer.name,
+                content=f"Longer message {i}",
+                token_count=limit // 2,  # Each message is half the limit
+                seq_in_session=len(messages) + i + 1,
+            )
+            db_session.add(message)
+            more_messages.append(message)
+
+        await db_session.commit()
+        for message in more_messages:
+            await db_session.refresh(message)
+
+        # Add queue items for new messages
+        for message in more_messages:
+            payload = create_queue_payload(
+                message=message,
+                task_type="representation",
+                observed=peer.name,
+                observer=peer.name,
+            )
+
+            queue_item = models.QueueItem(
+                session_id=session.id,
+                task_type="representation",
+                work_unit_key=rep_work_unit_key,  # Same work unit
+                payload=payload,
+                processed=False,
+                workspace_name=session.workspace_name,
+                message_id=message.id,
+            )
+            db_session.add(queue_item)
+
+        await db_session.commit()
+
+        # Now the work unit should be claimable (tokens exceed threshold)
+        claimed2 = await qm.get_and_claim_work_units()
+        assert rep_work_unit_key in claimed2
+
+    @pytest.mark.asyncio
+    async def test_forced_batching_single_large_message(
+        self,
+        db_session: AsyncSession,
+        sample_session_with_peers: tuple[models.Session, list[models.Peer]],
+        create_queue_payload: Callable[..., Any],
+    ) -> None:
+        """Test that a single message >= threshold is immediately claimable"""
+
+        session, peers = sample_session_with_peers
+        peer = peers[0]
+
+        limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
+
+        # Create a single message that exceeds the threshold
+        message = models.Message(
+            session_name=session.name,
+            workspace_name=session.workspace_name,
+            peer_name=peer.name,
+            content="A very long message",
+            token_count=limit + 1000,  # Exceeds threshold
+            seq_in_session=1,
+        )
+        db_session.add(message)
+        await db_session.commit()
+        await db_session.refresh(message)
+
+        # Create queue item
+        payload = create_queue_payload(
+            message=message,
+            task_type="representation",
+            observed=peer.name,
+            observer=peer.name,
+        )
+        work_unit_key = construct_work_unit_key(session.workspace_name, payload)
+
+        queue_item = models.QueueItem(
+            session_id=session.id,
+            task_type="representation",
+            work_unit_key=work_unit_key,
+            payload=payload,
+            processed=False,
+            workspace_name=session.workspace_name,
+            message_id=message.id,
+        )
+        db_session.add(queue_item)
+        await db_session.commit()
+
+        qm = QueueManager()
+
+        # Single large message should be immediately claimable
+        claimed = await qm.get_and_claim_work_units()
+        assert work_unit_key in claimed
+
+    @pytest.mark.asyncio
+    async def test_forced_batching_bypassed_for_summary(
+        self,
+        db_session: AsyncSession,
+        sample_session_with_peers: tuple[models.Session, list[models.Peer]],
+        create_queue_payload: Callable[..., Any],
+    ) -> None:
+        """Test that summary tasks are processed immediately regardless of tokens"""
+
+        session, peers = sample_session_with_peers
+        peer = peers[0]
+
+        # Create a message with very few tokens
+        message = models.Message(
+            session_name=session.name,
+            workspace_name=session.workspace_name,
+            peer_name=peer.name,
+            content="Short",
+            token_count=10,  # Way below threshold
+            seq_in_session=1,
+        )
+        db_session.add(message)
+        await db_session.commit()
+        await db_session.refresh(message)
+
+        # Create a SUMMARY queue item (not representation)
+        payload = create_queue_payload(
+            message=message,
+            task_type="summary",
+            message_seq_in_session=1,
+        )
+        work_unit_key = construct_work_unit_key(session.workspace_name, payload)
+
+        queue_item = models.QueueItem(
+            session_id=session.id,
+            task_type="summary",
+            work_unit_key=work_unit_key,
+            payload=payload,
+            processed=False,
+            workspace_name=session.workspace_name,
+            message_id=message.id,
+        )
+        db_session.add(queue_item)
+        await db_session.commit()
+
+        qm = QueueManager()
+
+        # Summary should be claimable regardless of token count
+        claimed = await qm.get_and_claim_work_units()
+        assert work_unit_key in claimed
+        assert work_unit_key.startswith("summary:")
+
+    @pytest.mark.asyncio
+    async def test_forced_batching_exact_threshold(
+        self,
+        db_session: AsyncSession,
+        sample_session_with_peers: tuple[models.Session, list[models.Peer]],
+        create_queue_payload: Callable[..., Any],
+    ) -> None:
+        """Test that work units exactly at the threshold are claimable"""
+
+        session, peers = sample_session_with_peers
+        peer = peers[0]
+
+        limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
+
+        # Create messages that sum to exactly the threshold
+        token_counts = [limit // 2, limit // 2]
+
+        messages: list[models.Message] = []
+        for i, token_count in enumerate(token_counts):
+            message = models.Message(
+                session_name=session.name,
+                workspace_name=session.workspace_name,
+                peer_name=peer.name,
+                content=f"Message {i}",
+                token_count=token_count,
+                seq_in_session=i + 1,
+            )
+            db_session.add(message)
+            messages.append(message)
+
+        await db_session.commit()
+        for message in messages:
+            await db_session.refresh(message)
+
+        # Create queue items
+        work_unit_key = None
+        for message in messages:
+            payload = create_queue_payload(
+                message=message,
+                task_type="representation",
+                observed=peer.name,
+                observer=peer.name,
+            )
+            if work_unit_key is None:
+                work_unit_key = construct_work_unit_key(session.workspace_name, payload)
+
+            queue_item = models.QueueItem(
+                session_id=session.id,
+                task_type="representation",
+                work_unit_key=work_unit_key,
+                payload=payload,
+                processed=False,
+                workspace_name=session.workspace_name,
+                message_id=message.id,
+            )
+            db_session.add(queue_item)
+
+        await db_session.commit()
+
+        qm = QueueManager()
+
+        # Work unit exactly at threshold should be claimable
+        claimed = await qm.get_and_claim_work_units()
+        assert work_unit_key is not None
+        assert work_unit_key in claimed

tests/dreamer/test_dream_scheduler.py

@@ -17,7 +17,9 @@ def dream_scheduler():
     DreamScheduler.reset_singleton()
     scheduler = DreamScheduler()
     set_dream_scheduler(scheduler)
-    yield scheduler
+    # Patch DREAM.ENABLED to True so tests work regardless of local config
+    with patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED", True):
+        yield scheduler
     # Cleanup
     DreamScheduler.reset_singleton()