feat: codify queue columns (#254)

* feat: codify queue columns

* fix: batch with python loop control

* fix: cleanup merge

* fix: down revision

* fix: batch delete in migration

* feat: only run alembic tests for changed migration / test (#264)

* feat: only run alembic tests for changed migration / test
* fix: Run full test suite if alembic testing infra changes

* feat: codify times_derived + level on Document (#260)

* feat: codify times_derived + level on Document
* fix: CR comment

* fix: CodeRabbit comments

* fix: batch migrations; move types; remove fields from payload

* fix: rm duplicate table args

* fix: add messages.id FK

.pre-commit-config.yaml

@@ -74,11 +74,11 @@ repos:
       # Run Alembic tests only when migrations change
       - id: pytest-alembic
         name: pytest (alembic migrations)
-        entry: uv run pytest tests/alembic/
+        entry: uv run python scripts/run_alembic_tests.py
         language: system
-        files: ^(migrations/.*\.py|tests/alembic/.*\.py)$
+        files: ^(migrations/versions/.*\.py|tests/alembic/.*\.py)$
         stages: [pre-push]
-        pass_filenames: false
+        pass_filenames: true
         require_serial: true
 
       # Ensure each alembic migration revision has a corresponding test file

migrations/versions/b8183c5ffb48_codify_document_level_and_times_derived.py

@@ -0,0 +1,164 @@
+"""codify_document_level_and_times_derived
+
+Revision ID: b8183c5ffb48
+Revises: ec8f94139b02
+Create Date: 2025-10-31 12:48:54.597269
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import text
+
+from migrations.utils import column_exists, constraint_exists, get_schema
+
+# revision identifiers, used by Alembic.
+revision: str = "b8183c5ffb48"
+down_revision: str | None = "ec8f94139b02"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+schema = get_schema()
+
+
+def upgrade() -> None:
+    """Codify level and times_derived from internal_metadata into explicit columns."""
+    connection = op.get_bind()
+    inspector = sa.inspect(connection)
+
+    # Step 1: Add level column (nullable initially)
+    if not column_exists("documents", "level", inspector):
+        op.add_column(
+            "documents",
+            sa.Column(
+                "level",
+                sa.TEXT(),
+                nullable=True,
+            ),
+            schema=schema,
+        )
+
+    # Step 2: Add times_derived column (nullable initially)
+    if not column_exists("documents", "times_derived", inspector):
+        op.add_column(
+            "documents",
+            sa.Column("times_derived", sa.Integer(), nullable=True),
+            schema=schema,
+        )
+
+    # Step 3: Populate level and times_derived from internal_metadata in batches
+    # Default to 'explicit' for level and 1 for times_derived if not present in metadata
+    batch_size = 5000
+    while True:
+        result = connection.execute(
+            text(
+                f"""
+                WITH batch AS (
+                    SELECT id
+                    FROM {schema}.documents
+                    WHERE level IS NULL OR times_derived IS NULL
+                    LIMIT :batch_size
+                )
+                UPDATE {schema}.documents d
+                SET
+                    level = COALESCE(
+                        d.internal_metadata->>'level',
+                        'explicit'
+                    ),
+                    times_derived = COALESCE(
+                        (d.internal_metadata->>'times_derived')::integer,
+                        1
+                    )
+                FROM batch
+                WHERE d.id = batch.id
+            """
+            ),
+            {"batch_size": batch_size},
+        )
+        if result.rowcount == 0:
+            break
+
+    # Step 4: Make level NOT NULL with server default
+    op.alter_column(
+        "documents",
+        "level",
+        nullable=False,
+        server_default=text("'explicit'"),
+        schema=schema,
+    )
+
+    # Step 5: Make times_derived NOT NULL with server default
+    op.alter_column(
+        "documents",
+        "times_derived",
+        nullable=False,
+        server_default=text("1"),
+        schema=schema,
+    )
+
+    # Step 6: Add CHECK constraint for level
+    if not constraint_exists("documents", "level_valid", "check", inspector):
+        op.create_check_constraint(
+            "level_valid",
+            "documents",
+            "level IN ('explicit', 'deductive')",
+            schema=schema,
+        )
+
+
+def downgrade() -> None:
+    """Restore level and times_derived to internal_metadata."""
+    connection = op.get_bind()
+    inspector = sa.inspect(connection)
+
+    # Step 1: Drop CHECK constraint for level
+    if constraint_exists("documents", "level_valid", "check", inspector):
+        op.drop_constraint(
+            "level_valid",
+            "documents",
+            type_="check",
+            schema=schema,
+        )
+
+    # Step 2: Copy level and times_derived back to internal_metadata in batches (optional, for safety)
+    batch_size = 5000
+    while True:
+        result = connection.execute(
+            text(
+                f"""
+                WITH batch AS (
+                    SELECT id
+                    FROM {schema}.documents
+                    WHERE internal_metadata IS NULL
+                       OR NOT (internal_metadata ? 'level')
+                       OR NOT (internal_metadata ? 'times_derived')
+                    LIMIT :batch_size
+                )
+                UPDATE {schema}.documents d
+                SET internal_metadata = jsonb_set(
+                    jsonb_set(
+                        COALESCE(d.internal_metadata, '{{}}'::jsonb),
+                        '{{level}}',
+                        to_jsonb(d.level)
+                    ),
+                    '{{times_derived}}',
+                    to_jsonb(d.times_derived)
+                )
+                FROM batch
+                WHERE d.id = batch.id
+            """
+            ),
+            {"batch_size": batch_size},
+        )
+        if result.rowcount == 0:
+            break
+
+    # Step 3: Drop the level column
+    if column_exists("documents", "level", inspector):
+        op.drop_column("documents", "level", schema=schema)
+
+    # Step 4: Drop the times_derived column
+    if column_exists("documents", "times_derived", inspector):
+        op.drop_column("documents", "times_derived", schema=schema)

migrations/versions/ec8f94139b02_codify_workspace_name_and_message_id_in_.py

@@ -0,0 +1,252 @@
+"""codify workspace_name and message_id in queue table
+
+Revision ID: ec8f94139b02
+Revises: e9b705f9adf9
+Create Date: 2025-10-28 17:39:51.778665
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+from migrations.utils import column_exists, fk_exists, get_schema, index_exists
+
+# revision identifiers, used by Alembic.
+revision: str = "ec8f94139b02"
+down_revision: str | None = "e9b705f9adf9"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+schema = get_schema()
+
+
+def upgrade() -> None:
+    # Step 1: Drop rows where workspace_name is NULL in payload
+    # This removes invalid/corrupted queue items
+    conn = op.get_bind()
+    batch_size = 10000
+    while True:
+        result = conn.execute(
+            sa.text(
+                f"""
+                DELETE FROM "{schema}".queue
+                WHERE id IN (
+                    SELECT id FROM "{schema}".queue
+                    WHERE payload->>'workspace_name' IS NULL
+                    LIMIT :batch_size
+                )
+                """
+            ),
+            {"batch_size": batch_size},
+        )
+        if result.rowcount == 0:
+            break
+
+    # Step 2: Add workspace_name column (nullable initially for backfill)
+    op.add_column(
+        "queue",
+        sa.Column("workspace_name", sa.TEXT(), nullable=True),
+        schema=schema,
+    )
+
+    # Step 3: Add message_id column (nullable, as not all tasks have message_id)
+    op.add_column(
+        "queue",
+        sa.Column("message_id", sa.BigInteger(), nullable=True),
+        schema=schema,
+    )
+
+    # Step 4: Backfill workspace_name and message_id from payload in batches
+    batch_size = 5000
+    while True:
+        result = conn.execute(
+            sa.text(
+                f"""
+                WITH batch AS (
+                    SELECT id
+                    FROM "{schema}".queue
+                    WHERE workspace_name IS NULL
+                       OR (message_id IS NULL AND payload ? 'message_id' AND payload->>'message_id' IS NOT NULL)
+                    LIMIT :batch_size
+                )
+                UPDATE "{schema}".queue q
+                SET
+                    workspace_name = COALESCE(q.workspace_name, q.payload->>'workspace_name'),
+                    message_id = COALESCE(
+                        q.message_id,
+                        CASE
+                            WHEN q.payload ? 'message_id' AND q.payload->>'message_id' IS NOT NULL
+                            THEN (q.payload->>'message_id')::bigint
+                        END
+                    )
+                FROM batch
+                WHERE q.id = batch.id
+                """
+            ),
+            {"batch_size": batch_size},
+        )
+        if result.rowcount == 0:
+            break
+
+    # Step 5: Remove workspace_name and message_id from JSONB payloads in batches
+    while True:
+        result = conn.execute(
+            sa.text(
+                f"""
+                WITH batch AS (
+                    SELECT id
+                    FROM "{schema}".queue
+                    WHERE payload ? 'workspace_name' OR payload ? 'message_id'
+                    LIMIT :batch_size
+                )
+                UPDATE "{schema}".queue q
+                SET payload = q.payload - 'workspace_name' - 'message_id'
+                FROM batch
+                WHERE q.id = batch.id
+                """
+            ),
+            {"batch_size": batch_size},
+        )
+        if result.rowcount == 0:
+            break
+
+    # Step 6: Make workspace_name non-nullable
+    op.alter_column("queue", "workspace_name", nullable=False, schema=schema)
+
+    # Step 7: Add foreign key constraint on workspace_name -> workspaces.name
+    op.create_foreign_key(
+        "fk_queue_workspace_name",
+        "queue",
+        "workspaces",
+        ["workspace_name"],
+        ["name"],
+        source_schema=schema,
+        referent_schema=schema,
+    )
+
+    # Step 7b: Add foreign key constraint on message_id -> messages.id
+    op.create_foreign_key(
+        "fk_queue_message_id",
+        "queue",
+        "messages",
+        ["message_id"],
+        ["id"],
+        source_schema=schema,
+        referent_schema=schema,
+    )
+
+    # Step 8: Add index on workspace_name (for FK performance and filtering)
+    op.create_index(
+        op.f("ix_queue_workspace_name"),
+        "queue",
+        ["workspace_name"],
+        unique=False,
+        schema=schema,
+    )
+
+    # Step 9: Add partial index on message_id WHERE message_id IS NOT NULL
+    # This optimizes JOINs with the messages table
+    op.create_index(
+        "ix_queue_message_id_not_null",
+        "queue",
+        ["message_id"],
+        unique=False,
+        schema=schema,
+        postgresql_where=sa.text("message_id IS NOT NULL"),
+    )
+
+    # Step 10: Add composite index on (workspace_name, processed)
+    # This optimizes queries that filter unprocessed items by workspace
+    op.create_index(
+        "ix_queue_workspace_name_processed",
+        "queue",
+        ["workspace_name", "processed"],
+        unique=False,
+        schema=schema,
+    )
+
+    # Step 11: Add composite index on (work_unit_key, processed, id)
+    # This is critical for the hot path: "get next unprocessed item for this work unit"
+    # Covers: WHERE work_unit_key = ? AND NOT processed ORDER BY id
+    op.create_index(
+        "ix_queue_work_unit_key_processed_id",
+        "queue",
+        ["work_unit_key", "processed", "id"],
+        unique=False,
+        schema=schema,
+    )
+
+
+def downgrade() -> None:
+    inspector = sa.inspect(op.get_bind())
+
+    # Drop indexes
+    if index_exists("queue", "ix_queue_work_unit_key_processed_id", inspector):
+        op.drop_index(
+            "ix_queue_work_unit_key_processed_id", table_name="queue", schema=schema
+        )
+
+    if index_exists("queue", "ix_queue_workspace_name_processed", inspector):
+        op.drop_index(
+            "ix_queue_workspace_name_processed", table_name="queue", schema=schema
+        )
+
+    if index_exists("queue", "ix_queue_message_id_not_null", inspector):
+        op.drop_index("ix_queue_message_id_not_null", table_name="queue", schema=schema)
+
+    if index_exists("queue", "ix_queue_workspace_name", inspector):
+        op.drop_index(
+            op.f("ix_queue_workspace_name"), table_name="queue", schema=schema
+        )
+
+    # Drop foreign key constraints
+    if fk_exists("queue", "fk_queue_message_id", inspector):
+        op.drop_constraint("fk_queue_message_id", "queue", schema=schema)
+
+    if fk_exists("queue", "fk_queue_workspace_name", inspector):
+        op.drop_constraint("fk_queue_workspace_name", "queue", schema=schema)
+
+    # Restore workspace_name and message_id to payload in batches
+    conn = op.get_bind()
+    batch_size = 5000
+
+    if column_exists("queue", "workspace_name", inspector) or column_exists(
+        "queue", "message_id", inspector
+    ):
+        while True:
+            result = conn.execute(
+                sa.text(
+                    f"""
+                    WITH batch AS (
+                        SELECT id
+                        FROM "{schema}".queue
+                        WHERE (workspace_name IS NOT NULL AND NOT (payload ? 'workspace_name'))
+                           OR (message_id IS NOT NULL AND NOT (payload ? 'message_id'))
+                        LIMIT :batch_size
+                    )
+                    UPDATE "{schema}".queue q
+                    SET payload = q.payload
+                        || CASE WHEN q.workspace_name IS NOT NULL AND NOT (q.payload ? 'workspace_name')
+                                THEN jsonb_build_object('workspace_name', q.workspace_name)
+                                ELSE '{{}}'::jsonb END
+                        || CASE WHEN q.message_id IS NOT NULL AND NOT (q.payload ? 'message_id')
+                                THEN jsonb_build_object('message_id', q.message_id)
+                                ELSE '{{}}'::jsonb END
+                    FROM batch
+                    WHERE q.id = batch.id
+                    """
+                ),
+                {"batch_size": batch_size},
+            )
+            if result.rowcount == 0:
+                break
+
+    # Drop columns
+    if column_exists("queue", "message_id", inspector):
+        op.drop_column("queue", "message_id", schema=schema)
+
+    if column_exists("queue", "workspace_name", inspector):
+        op.drop_column("queue", "workspace_name", schema=schema)

scripts/run_alembic_tests.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""
+Run alembic tests selectively based on changed files.
+
+Note that this script is actively used within our precommit hooks and should not be removed.
+If this script is moved, the corresponding precommit hook will need to be updated.
+
+This script determines which specific alembic tests to run based on:
+1. If a test file changed, run the test for that revision
+2. If a migration file changed, run the corresponding test
+
+The alembic test system uses a parameterized test in test_pipeline.py that runs
+for each revision. We filter these tests using pytest's -k flag with the revision IDs.
+
+Usage: python scripts/run_alembic_tests.py <file1> <file2> ...
+"""
+
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+
+def extract_revision_id(filepath: Path) -> str | None:
+    """Extract the revision ID from a migration or test filename.
+
+    Migration files: {revision_id}_{description}.py
+    Test files: test_{revision_id}_{description}.py
+
+    Returns the revision_id (e.g., "05486ce795d5") or None if not found.
+    """
+    filename = filepath.name
+
+    # Remove .py extension
+    if not filename.endswith(".py"):
+        return None
+
+    filename = filename[:-3]
+
+    # Remove test_ prefix if present
+    if filename.startswith("test_"):
+        filename = filename[5:]
+
+    # Extract revision ID (first part before underscore)
+    # Revision IDs are typically 12 characters of hex
+    match = re.match(r"^([a-f0-9]{12})_", filename)
+    if match:
+        return match.group(1)
+
+    return None
+
+
+def main():
+    if len(sys.argv) < 2:
+        print("No files to check, skipping alembic tests")
+        sys.exit(0)
+
+    changed_files = [Path(f) for f in sys.argv[1:]]
+
+    # Paths
+    repo_root = Path(__file__).parent.parent
+    migrations_dir = repo_root / "migrations" / "versions"
+    tests_dir = repo_root / "tests" / "alembic" / "revisions"
+    alembic_tests_dir = repo_root / "tests" / "alembic"
+
+    # Collect revision IDs to test
+    revision_ids: set[str] = set()
+    run_full_suite = False
+
+    for filepath in changed_files:
+        filepath = Path(filepath).resolve()
+
+        # Check if file is under tests/alembic (including subdirectories)
+        # If it's not a revision-specific test file, run full suite
+        if (
+            filepath.parent == alembic_tests_dir
+            or alembic_tests_dir in filepath.parents
+        ) and not (filepath.parent == tests_dir and filepath.name.startswith("test_")):
+            run_full_suite = True
+            print(
+                f"Infrastructure file changed: {filepath.name} -> will run full test suite"
+            )
+            continue
+
+        # Case 1: Test file changed - extract its revision ID
+        if filepath.parent == tests_dir and filepath.name.startswith("test_"):
+            revision_id = extract_revision_id(filepath)
+            if revision_id:
+                revision_ids.add(revision_id)
+                print(
+                    f"Test file changed: {filepath.name} -> testing revision {revision_id}"
+                )
+
+        # Case 2: Migration file changed - extract its revision ID
+        elif filepath.parent == migrations_dir:
+            revision_id = extract_revision_id(filepath)
+            if revision_id:
+                revision_ids.add(revision_id)
+                print(
+                    f"Migration changed: {filepath.name} -> testing revision {revision_id}"
+                )
+
+    if run_full_suite:
+        # Run full test suite without -k filter
+        print("\nRunning full alembic test suite due to infrastructure file changes\n")
+        cmd = [
+            "uv",
+            "run",
+            "pytest",
+            "tests/alembic/test_pipeline.py",
+        ]
+    elif revision_ids:
+        # Build a -k expression to filter tests by revision ID
+        # pytest -k "rev1 or rev2 or rev3"
+        k_expression = " or ".join(sorted(revision_ids))
+
+        print(
+            f"\nRunning tests for {len(revision_ids)} revision(s): {', '.join(sorted(revision_ids))}"
+        )
+        print()
+
+        # Run pytest on test_pipeline.py with -k filter
+        cmd = [
+            "uv",
+            "run",
+            "pytest",
+            "tests/alembic/test_pipeline.py",
+            "-k",
+            k_expression,
+        ]
+    else:
+        print("No alembic tests to run")
+        sys.exit(0)
+
+    result = subprocess.run(cmd, cwd=repo_root)
+
+    sys.exit(result.returncode)
+
+
+if __name__ == "__main__":
+    main()

src/crud/document.py

@@ -127,6 +127,8 @@ async def create_documents(
                     observer=observer,
                     observed=observed,
                     content=doc.content,
+                    level=doc.level,
+                    times_derived=doc.times_derived,
                     internal_metadata=metadata_dict,
                     embedding=doc.embedding,
                     session_name=doc.session_name,

src/crud/representation.py

@@ -145,7 +145,6 @@ class RepresentationManager:
 
             metadata: schemas.DocumentMetadata = schemas.DocumentMetadata(
                 message_ids=[message_id_range],
-                level=obs_level,
                 premises=obs_premises,
                 message_created_at=format_datetime_utc(message_created_at),
             )
@@ -154,6 +153,7 @@ class RepresentationManager:
                 schemas.DocumentCreate(
                     content=obs_content,
                     session_name=session_name,
+                    level=obs_level,
                     metadata=metadata,
                     embedding=embedding,
                 )
@@ -397,7 +397,7 @@ class RepresentationManager:
                 models.Document.observer == self.observer,
                 models.Document.observed == self.observed,
             )
-            .order_by(models.Document.internal_metadata["times_derived"].desc())
+            .order_by(models.Document.times_derived.desc())
         )
 
         result = await db.execute(stmt)
@@ -454,7 +454,7 @@ class RepresentationManager:
         conditions: list[dict[str, Any]] = []
 
         if level:
-            conditions.append({"internal_metadata": {"level": level}})
+            conditions.append({"level": level})
 
         if not conditions:
             return {}

src/deriver/consumer.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Any
 
 import sentry_sdk
 from pydantic import ValidationError
@@ -26,8 +25,12 @@ logging.getLogger("sqlalchemy.engine.Engine").disabled = True
 console = Console(markup=True)
 
 
-async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
+async def process_item(queue_item: models.QueueItem) -> None:
     """Process a single item from the queue."""
+    task_type = queue_item.task_type
+    queue_payload = queue_item.payload
+    workspace_name = queue_item.workspace_name
+
     if task_type == "webhook":
         try:
             validated = WebhookPayload(**queue_payload)
@@ -39,7 +42,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
             )
             raise ValueError(f"Invalid payload structure: {str(e)}") from e
         async with tracked_db() as db:
-            await webhook_delivery.deliver_webhook(db, validated)
+            await webhook_delivery.deliver_webhook(db, validated, workspace_name)
 
     elif task_type == "summary":
         try:
@@ -52,17 +55,21 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
             )
             raise ValueError(f"Invalid payload structure: {str(e)}") from e
 
+        # Use workspace_name and message_id from QueueItem columns
+        message_id = queue_item.message_id
+
+        if message_id is None:
+            raise ValueError("Summary tasks require a message_id")
+
         message_public_id = validated.message_public_id
         if not message_public_id:
-            logger.debug(
-                "Fetching message public ID for message %s", validated.message_id
-            )
+            logger.debug("Fetching message public ID for message %s", message_id)
             async with tracked_db(operation_name="summary_fallback") as db:
                 stmt = (
                     select(models.Message)
-                    .where(models.Message.workspace_name == validated.workspace_name)
+                    .where(models.Message.workspace_name == workspace_name)
                     .where(models.Message.session_name == validated.session_name)
-                    .where(models.Message.id == validated.message_id)
+                    .where(models.Message.id == message_id)
                 )
                 result = await db.execute(stmt)
 
@@ -70,22 +77,20 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
                 if message is None:
                     logger.error(
                         "Failed to fetch message with ID %s for process_summary_task",
-                        validated.message_id,
+                        message_id,
                     )
                     return
                 message_public_id = message.public_id
 
         with sentry_sdk.start_transaction(name="process_summary_task", op="deriver"):
             await summarizer.summarize_if_needed(
-                validated.workspace_name,
+                workspace_name,
                 validated.session_name,
-                validated.message_id,
+                message_id,
                 validated.message_seq_in_session,
                 message_public_id,
             )
-            log_performance_metrics(
-                "summary", f"{validated.workspace_name}_{validated.message_id}"
-            )
+            log_performance_metrics("summary", f"{workspace_name}_{message_id}")
 
     elif task_type == "dream":
         with sentry_sdk.start_transaction(name="process_dream_task", op="deriver"):
@@ -98,7 +103,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
                     queue_payload,
                 )
                 raise ValueError(f"Invalid payload structure: {str(e)}") from e
-            await process_dream(validated)
+            await process_dream(validated, workspace_name)
     else:
         raise ValueError(f"Invalid task type: {task_type}")
 
@@ -109,29 +114,21 @@ async def process_representation_batch(
     observer: str | None,
     observed: str | None,
 ) -> None:
-    """Validate incoming queue payloads and dispatch to the appropriate handler.
-
-    This function centralizes payload validation using a simple mapping from
-    task type to Pydantic model. After validation, routes the request to
-    the correct processor without repeating type checks elsewhere.
-
+    """Prepares and processes a batch of messages for representation tasks.
     Args:
-        task_type: The type of task to process
-        queue_payloads: List of payload dictionaries to process
-        observed (optional): For representation tasks, the observed from work_unit_key
-                     to identify which messages should be focused on
-        observer (optional): For representation tasks, the observer from work_unit_key
-                     to identify which messages should be focused on
+        messages: List of messages to process
+        observer: The observer of the messages
+        observed: The observed of the messages
     """
     if not messages or not messages[0]:
-        logger.debug("process_representation_batch received no payloads")
+        logger.debug("process_representation_batch received no messages")
         return
 
     if observed is None or observer is None:
         raise ValueError("observed and observer are required for representation tasks")
 
     logger.debug(
-        "process_representation_batch received %s payloads",
+        "process_representation_batch received %s messages",
         len(messages),
     )
 

src/deriver/enqueue.py

@@ -11,7 +11,7 @@ from src.dreamer.dream_scheduler import get_affected_dream_keys, get_dream_sched
 from src.exceptions import ValidationException
 from src.models import QueueItem
 from src.utils.queue_payload import create_payload
-from src.utils.work_unit import get_work_unit_key
+from src.utils.work_unit import construct_work_unit_key
 
 logger = logging.getLogger(__name__)
 
@@ -159,8 +159,16 @@ def create_representation_record(
         session_id: Optional session ID
 
     Returns:
-        Queue record dictionary
+        Queue record dictionary with workspace_name and message_id as separate fields
     """
+    workspace_name = message.get("workspace_name")
+    message_id = message.get("message_id")
+
+    if not isinstance(workspace_name, str):
+        raise TypeError("workspace_name is required and must be a string")
+    if not isinstance(message_id, int):
+        raise TypeError("message_id is required and must be an integer")
+
     processed_payload = create_payload(
         message=message,
         task_type="representation",
@@ -168,10 +176,12 @@ def create_representation_record(
         observed=observed,
     )
     return {
-        "work_unit_key": get_work_unit_key(processed_payload),
+        "work_unit_key": construct_work_unit_key(workspace_name, processed_payload),
         "payload": processed_payload,
         "session_id": session_id,
         "task_type": "representation",
+        "workspace_name": workspace_name,
+        "message_id": message_id,
     }
 
 
@@ -185,23 +195,32 @@ def create_summary_record(
 
     Args:
         message: The message payload
-        observed: Name of the sender
-        observer: Name of the target
         session_id: Session ID
+        message_seq_in_session: The sequence number of the message in the session
 
     Returns:
-        Queue record dictionary
+        Queue record dictionary with workspace_name and message_id as separate fields
     """
+    workspace_name = message.get("workspace_name")
+    message_id = message.get("message_id")
+
+    if not isinstance(workspace_name, str):
+        raise ValueError("workspace_name is required and must be a string")
+    if not isinstance(message_id, int):
+        raise ValueError("message_id is required and must be an integer")
+
     processed_payload = create_payload(
         message=message,
         task_type="summary",
         message_seq_in_session=message_seq_in_session,
     )
     return {
-        "work_unit_key": get_work_unit_key(processed_payload),
+        "work_unit_key": construct_work_unit_key(workspace_name, processed_payload),
         "payload": processed_payload,
         "session_id": session_id,
         "task_type": "summary",
+        "workspace_name": workspace_name,
+        "message_id": message_id,
     }
 
 

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 BigInteger, and_, delete, select, update
+from sqlalchemy import and_, delete, select, update
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.sql import func
@@ -464,9 +464,7 @@ class QueueManager:
                                 break
 
                             try:
-                                await process_item(
-                                    work_unit.task_type, queue_item.payload
-                                )
+                                await process_item(queue_item)
                                 await self.mark_queue_items_as_processed(
                                     [queue_item], work_unit_key
                                 )
@@ -615,8 +613,7 @@ class QueueManager:
                 .select_from(models.QueueItem)
                 .join(
                     models.Message,
-                    func.cast(models.QueueItem.payload["message_id"].astext, BigInteger)
-                    == models.Message.id,
+                    models.QueueItem.message_id == models.Message.id,
                 )
                 .where(~models.QueueItem.processed)
                 .where(models.Message.session_name == parsed_key.session_name)
@@ -662,10 +659,7 @@ class QueueManager:
                     and_(
                         models.QueueItem.work_unit_key == work_unit_key,
                         ~models.QueueItem.processed,
-                        func.cast(
-                            models.QueueItem.payload["message_id"].astext, BigInteger
-                        )
-                        == models.Message.id,
+                        models.QueueItem.message_id == models.Message.id,
                     ),
                 )
                 .where(allowed_condition)
@@ -690,7 +684,11 @@ class QueueManager:
 
             if items_to_process:
                 max_queue_item_message_id = max(
-                    [qi.payload["message_id"] for qi in items_to_process]
+                    [
+                        qi.message_id
+                        for qi in items_to_process
+                        if qi.message_id is not None
+                    ]
                 )
                 messages_context = [  # remove any messages that are after the last message_id from queue items
                     m for m in messages_context if m.id <= max_queue_item_message_id

src/dreamer/dream_scheduler.py

@@ -12,7 +12,7 @@ from src import models
 from src.config import settings
 from src.dependencies import tracked_db
 from src.utils.queue_payload import create_dream_payload
-from src.utils.work_unit import get_work_unit_key, parse_work_unit_key
+from src.utils.work_unit import construct_work_unit_key, parse_work_unit_key
 
 logger = getLogger(__name__)
 
@@ -48,13 +48,13 @@ def get_affected_dream_keys(message: dict[str, Any]) -> list[str]:
         return []
 
     # Generate dream work unit key for this peer's collection
-    dream_key = get_work_unit_key(
+    dream_key = construct_work_unit_key(
+        workspace_name,
         {
             "task_type": "dream",
-            "workspace_name": workspace_name,
             "observer": peer_name,
             "observed": peer_name,
-        }
+        },
     )
 
     return [dream_key]
@@ -193,7 +193,6 @@ class DreamScheduler:
     ) -> None:
         """Execute the dream by enqueueing it and updating collection metadata."""
         dream_payload = create_dream_payload(
-            workspace_name=workspace_name,
             dream_type="consolidate",
             observer=observer,
             observed=observed,
@@ -205,6 +204,8 @@ class DreamScheduler:
                 "payload": dream_payload,
                 "session_id": None,
                 "task_type": "dream",
+                "workspace_name": workspace_name,
+                "message_id": None,  # Dreams don't have a message_id
             }
 
             await db.execute(insert(models.QueueItem), [dream_record])
@@ -325,13 +326,13 @@ async def check_and_schedule_dream(
 
         dream_scheduler = get_dream_scheduler()
         if dream_scheduler:
-            collection_work_unit_key = get_work_unit_key(
+            collection_work_unit_key = construct_work_unit_key(
+                collection.workspace_name,
                 {
                     "task_type": "dream",
-                    "workspace_name": collection.workspace_name,
                     "observer": collection.observer,
                     "observed": collection.observed,
-                }
+                },
             )
 
             await dream_scheduler.schedule_dream(

src/dreamer/dreamer.py

@@ -24,6 +24,7 @@ logger = logging.getLogger(__name__)
 @sentry_sdk.trace
 async def process_dream(
     payload: DreamPayload,
+    workspace_name: str,
 ) -> None:
     """
     Process a dream task by performing collection maintenance operations.
@@ -32,12 +33,12 @@ async def process_dream(
         payload: The dream task payload containing workspace, peer, and dream type information
     """
     logger.info(
-        f"Processing dream task: {payload.dream_type} for {payload.workspace_name}/{payload.observer}/{payload.observed}"
+        f"Processing dream task: {payload.dream_type} for {workspace_name}/{payload.observer}/{payload.observed}"
     )
 
     try:
         if payload.dream_type == "consolidate":
-            await _process_consolidate_dream(payload)
+            await _process_consolidate_dream(payload, workspace_name)
         ## TODO other dream types
 
     except Exception as e:
@@ -50,7 +51,9 @@ async def process_dream(
         # Don't re-raise - we want to mark the dream task as processed even if it fails
 
 
-async def _process_consolidate_dream(payload: DreamPayload) -> None:
+async def _process_consolidate_dream(
+    payload: DreamPayload, workspace_name: str
+) -> None:
     """
     Process a consolidation dream task.
 
@@ -63,7 +66,7 @@ async def _process_consolidate_dream(payload: DreamPayload) -> None:
     logger.info(
         f"""
 (っ- ‸ - ς)ᶻ z 𐰁 ᶻ z 𐰁 ᶻ z 𐰁\n
-DREAM: consolidating documents for {payload.workspace_name}/{payload.observer}/{payload.observed}\n
+DREAM: consolidating documents for {workspace_name}/{payload.observer}/{payload.observed}\n
 𐰁 z ᶻ 𐰁 z ᶻ 𐰁 z ᶻ(っ- ‸ - ς)"""
     )
 
@@ -71,7 +74,7 @@ DREAM: consolidating documents for {payload.workspace_name}/{payload.observer}/{
     async with tracked_db("dream_consolidate") as db:
         documents = await crud.get_all_documents(
             db,
-            payload.workspace_name,
+            workspace_name,
             observer=payload.observer,
             observed=payload.observed,
         )
@@ -88,7 +91,7 @@ DREAM: consolidating documents for {payload.workspace_name}/{payload.observer}/{
         for cluster in clusters:
             await _consolidate_cluster(
                 cluster,
-                payload.workspace_name,
+                workspace_name,
                 db,
                 observer=payload.observer,
                 observed=payload.observed,
@@ -118,9 +121,7 @@ async def _consolidate_cluster(
     logger.info("consolidated representation:\n%s", consolidated_representation)
 
     # TODO: less hacky preservation of times_derived
-    total_times_derived = sum(
-        doc.internal_metadata.get("times_derived", 1) for doc in cluster
-    )
+    total_times_derived = sum(doc.times_derived for doc in cluster)
 
     new_documents = [
         *consolidated_representation.explicit,
@@ -141,10 +142,8 @@ async def _consolidate_cluster(
         # NOTE: other kinds of observations here in the future
 
         metadata = schemas.DocumentMetadata(
-            times_derived=total_times_derived,
             message_ids=obs.message_ids,
             message_created_at=format_datetime_utc(obs.created_at),
-            level=level,
             premises=premises,
         )
 
@@ -154,6 +153,8 @@ async def _consolidate_cluster(
             schemas.DocumentCreate(
                 content=content,
                 session_name=obs.session_name,
+                level=level,
+                times_derived=total_times_derived,
                 metadata=metadata,
                 embedding=embedding,
             )

src/models.py

@@ -1,6 +1,6 @@
 import datetime
 from logging import getLogger
-from typing import Any, Literal, final
+from typing import Any, final
 
 from dotenv import load_dotenv
 from nanoid import generate as generate_nanoid
@@ -26,6 +26,8 @@ from sqlalchemy.orm.properties import MappedColumn
 from sqlalchemy.sql import func
 from typing_extensions import override
 
+from src.utils.types import DocumentLevel, TaskType
+
 from .db import Base
 
 load_dotenv(override=True)
@@ -382,6 +384,12 @@ class Document(Base):
         "internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
     )
     content: Mapped[str] = mapped_column(TEXT)
+    level: Mapped[DocumentLevel] = mapped_column(
+        TEXT, nullable=False, server_default="explicit"
+    )
+    times_derived: Mapped[int] = mapped_column(
+        Integer, nullable=False, server_default=text("1")
+    )
     embedding: MappedColumn[Any] = mapped_column(Vector(1536))
     created_at: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True), server_default=func.now()
@@ -397,6 +405,7 @@ class Document(Base):
         CheckConstraint("length(id) = 21", name="id_length"),
         CheckConstraint("length(content) <= 65535", name="content_length"),
         CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
+        CheckConstraint("level IN ('explicit', 'deductive')", name="level_valid"),
         # Composite foreign key constraint for collections
         ForeignKeyConstraint(
             ["observer", "observed", "workspace_name"],
@@ -439,9 +448,6 @@ class Document(Base):
     )
 
 
-TaskType = Literal["webhook", "summary", "representation", "dream"]
-
-
 @final
 class QueueItem(Base):
     __tablename__: str = "queue"
@@ -460,14 +466,36 @@ class QueueItem(Base):
     created_at: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True), server_default=func.now()
     )
+    workspace_name: Mapped[str] = mapped_column(
+        ForeignKey("workspaces.name"), nullable=False
+    )
+    message_id: Mapped[int | None] = mapped_column(
+        BigInteger, ForeignKey("messages.id"), nullable=True
+    )
 
     __table_args__ = (
         Index("ix_queue_created_at", "created_at"),
         Index("ix_queue_session_id", "session_id"),
+        Index(
+            "ix_queue_workspace_name",
+            "workspace_name",
+        ),
+        Index(
+            "ix_queue_message_id_not_null",
+            "message_id",
+            postgresql_where=text("message_id IS NOT NULL"),
+        ),
+        Index("ix_queue_workspace_name_processed", "workspace_name", "processed"),
+        Index(
+            "ix_queue_work_unit_key_processed_id",
+            "work_unit_key",
+            "processed",
+            "id",
+        ),
     )
 
     def __repr__(self) -> str:
-        return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed})"
+        return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed}, workspace_name={self.workspace_name}, message_id={self.message_id})"
 
 
 @final

src/schemas.py

@@ -288,20 +288,12 @@ class DocumentBase(BaseModel):
 
 
 class DocumentMetadata(BaseModel):
-    times_derived: int | None = Field(
-        default=None,
-        ge=1,
-        description="The number of times that a semantic duplicate document to this one has been derived",
-    )
     message_ids: list[tuple[int, int]] = Field(
         description="The ID range(s) of the messages that this document was derived from. Acts as a link to the primary source of the document. Note that as a document gets deduplicated, additional ranges will be added, because the same document could be derived from completely separate message ranges."
     )
     message_created_at: str = Field(
         description="The timestamp of the message that this document was derived from. Note that this is not the same as the created_at timestamp of the document. This timestamp is usually only saved with second-level precision."
     )
-    level: Literal["explicit", "deductive"] = Field(
-        description="The level of the document (explicit or deductive)"
-    )
     premises: list[str] | None = Field(
         default=None,
         description="The premises of the deduction -- only applicable for deductive observations",
@@ -313,6 +305,15 @@ class DocumentCreate(DocumentBase):
     session_name: str = Field(
         description="The session from which the document was derived"
     )
+    level: Literal["explicit", "deductive"] = Field(
+        default="explicit",
+        description="The level of the document (explicit or deductive)",
+    )
+    times_derived: int = Field(
+        default=1,
+        ge=1,
+        description="The number of times that a semantic duplicate document to this one has been derived",
+    )
     metadata: DocumentMetadata = Field()
     embedding: list[float] = Field()
 

src/utils/queue_payload.py

@@ -14,9 +14,7 @@ class RepresentationPayload(BasePayload):
     """Payload for representation tasks."""
 
     task_type: Literal["representation"] = "representation"
-    workspace_name: str
     session_name: str
-    message_id: int
     content: str
     observer: str
     observed: str
@@ -33,9 +31,7 @@ class SummaryPayload(BasePayload):
     """Payload for summary tasks."""
 
     task_type: Literal["summary"] = "summary"
-    workspace_name: str
     session_name: str
-    message_id: int
     message_seq_in_session: int
     # Optional for backward compatibility with older queue items
     message_public_id: str | None = None
@@ -45,7 +41,6 @@ class WebhookPayload(BasePayload):
     """Payload for webhook delivery tasks."""
 
     task_type: Literal["webhook"] = "webhook"
-    workspace_name: str
     event_type: str
     data: dict[str, Any]
 
@@ -54,35 +49,33 @@ class DreamPayload(BasePayload):
     """Payload for dream tasks."""
 
     task_type: Literal["dream"] = "dream"
-    workspace_name: str
     dream_type: Literal["consolidate"] = "consolidate"
     observer: str
     observed: str
 
 
 def create_webhook_payload(
-    workspace_name: str,
     event_type: str,
     data: dict[str, Any],
 ) -> dict[str, Any]:
-    return WebhookPayload(
-        workspace_name=workspace_name, event_type=event_type, data=data
-    ).model_dump(mode="json")
+    """Create a webhook payload."""
+    return WebhookPayload(event_type=event_type, data=data).model_dump(
+        mode="json", exclude_none=True
+    )
 
 
 def create_dream_payload(
-    workspace_name: str,
     dream_type: Literal["consolidate"] = "consolidate",
     *,
     observer: str,
     observed: str,
 ) -> dict[str, Any]:
+    """Create a dream payload."""
     return DreamPayload(
-        workspace_name=workspace_name,
         dream_type=dream_type,
         observer=observer,
         observed=observed,
-    ).model_dump(mode="json")
+    ).model_dump(mode="json", exclude_none=True)
 
 
 def create_payload(
@@ -96,6 +89,10 @@ def create_payload(
     """
     Create a processed payload from a message for queue processing.
 
+    Note: workspace_name and message_id are no longer included in the returned payload
+    as they are now stored in dedicated columns on the queue table. The caller is
+    responsible for extracting and passing these values separately.
+
     Args:
         message: The original message dictionary
         task_type: Type of task ('representation' or 'summary')
@@ -104,7 +101,7 @@ def create_payload(
         message_seq_in_session: Required for summary tasks, must be None for representation
 
     Returns:
-        Processed payload dictionary ready for queue processing
+        Processed payload dictionary ready for queue processing (without workspace_name and message_id)
 
     Raises:
         ValueError: If the payload doesn't match the expected schema
@@ -142,9 +139,7 @@ def create_payload(
 
             validated_payload = RepresentationPayload(
                 content=content,
-                workspace_name=workspace_name,
                 session_name=session_name,
-                message_id=message_id,
                 created_at=created_at,
                 observer=observer,
                 observed=observed,
@@ -161,16 +156,14 @@ def create_payload(
                 )
 
             validated_payload = SummaryPayload(
-                workspace_name=workspace_name,
                 session_name=session_name,
-                message_id=message_id,
                 message_seq_in_session=message_seq_in_session,
                 message_public_id=message_public_id,
             )
 
         # Convert back to dict for compatibility with JSON serialization
         # mode='json' ensures datetime is converted to ISO string
-        payload = validated_payload.model_dump(mode="json")
+        payload = validated_payload.model_dump(mode="json", exclude_none=True)
 
     except Exception as e:
         raise ValueError(f"Failed to create valid payload: {str(e)}") from e

src/utils/representation.py

@@ -271,7 +271,7 @@ class Representation(BaseModel):
                     session_name=doc.session_name,
                 )
                 for doc in documents
-                if doc.internal_metadata.get("level") == "explicit"
+                if doc.level == "explicit"
             ],
             deductive=[
                 DeductiveObservation(
@@ -284,7 +284,7 @@ class Representation(BaseModel):
                     premises=doc.internal_metadata.get("premises", []),
                 )
                 for doc in documents
-                if doc.internal_metadata.get("level") == "deductive"
+                if doc.level == "deductive"
             ],
         )
 

src/utils/types.py

@@ -1,3 +1,5 @@
 from typing import Literal
 
 SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom", "vllm"]
+TaskType = Literal["webhook", "summary", "representation", "dream"]
+DocumentLevel = Literal["explicit", "deductive"]

src/utils/work_unit.py

@@ -15,11 +15,14 @@ class ParsedWorkUnit(BaseModel):
     observed: str | None
 
 
-def get_work_unit_key(payload: dict[str, Any] | ParsedWorkUnit) -> str:
+def construct_work_unit_key(
+    workspace_name: str, payload: dict[str, Any] | ParsedWorkUnit
+) -> str:
     """
     Generate a work unit key for a given task type, workspace name, and event type.
 
     Args:
+        workspace_name: The name of the workspace the work unit belongs to
         payload: Dictionary containing work unit information
 
     Returns:
@@ -31,7 +34,6 @@ def get_work_unit_key(payload: dict[str, Any] | ParsedWorkUnit) -> str:
     if isinstance(payload, ParsedWorkUnit):
         payload = payload.model_dump()
 
-    workspace_name: str | None = payload.get("workspace_name")
     task_type: str | None = payload.get("task_type")
     if not workspace_name or not task_type:
         raise ValueError(

src/webhooks/events.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel
 from src.dependencies import tracked_db
 from src.models import QueueItem
 from src.utils.queue_payload import create_webhook_payload
-from src.utils.work_unit import get_work_unit_key
+from src.utils.work_unit import construct_work_unit_key
 
 logger = logging.getLogger(__name__)
 
@@ -51,23 +51,26 @@ async def publish_webhook_event(event: WebhookEvent) -> None:
         event: The webhook event to publish.
     """
     try:
+        # Note: workspace_name is no longer included in the payload
+        # It's stored directly on the queue item
         payload = create_webhook_payload(
-            workspace_name=event.workspace_id,
             event_type=event.type.value,
             data=event.model_dump(mode="json", exclude={"type"}),
         )
 
         async with tracked_db("publish_webhook_event") as db:
             queue_item = QueueItem(
-                work_unit_key=get_work_unit_key(
+                work_unit_key=construct_work_unit_key(
+                    event.workspace_id,
                     {
                         "task_type": "webhook",
-                        "workspace_name": event.workspace_id,
-                    }
+                    },
                 ),
                 payload=payload,
                 session_id=None,
                 task_type="webhook",
+                workspace_name=event.workspace_id,
+                message_id=None,  # Webhooks don't have a message_id
             )
             db.add(queue_item)
             await db.commit()

src/webhooks/webhook_delivery.py

@@ -15,16 +15,18 @@ from src.utils.queue_payload import WebhookPayload
 logger = logging.getLogger(__name__)
 
 
-async def deliver_webhook(db: AsyncSession, payload: WebhookPayload) -> None:
+async def deliver_webhook(
+    db: AsyncSession, payload: WebhookPayload, workspace_name: str
+) -> None:
     """
     Deliver a single webhook event to its configured endpoints.
     """
     async with httpx.AsyncClient(timeout=30.0) as client:
         try:
-            webhook_urls = await _get_webhook_urls(db, payload.workspace_name)
+            webhook_urls = await _get_webhook_urls(db, workspace_name)
             if not webhook_urls:
                 logger.debug(
-                    f"No webhook endpoints for workspace {payload.workspace_name}, skipping."
+                    f"No webhook endpoints for workspace {workspace_name}, skipping."
                 )
                 return
 
@@ -72,7 +74,7 @@ async def deliver_webhook(db: AsyncSession, payload: WebhookPayload) -> None:
                     )
 
         except httpx.RequestError:
-            logger.exception(f"Error sending webhook for {payload.workspace_name}.")
+            logger.exception(f"Error sending webhook for {workspace_name}.")
         except Exception:
             logger.exception("Unexpected error delivering webhook.")
 

tests/alembic/revisions/__init__.py

@@ -13,9 +13,11 @@ from . import (
     test_917195d9b5e9_add_messageembedding_table,
     test_a1b2c3d4e5f6_initial_schema,
     test_b765d82110bd_change_metamessages_to_user_level_with_,
+    test_b8183c5ffb48_codify_document_level_and_times_derived,
     test_bb6fb3a7a643_add_message_seq_in_session_column,
     test_c3828084f472_add_indexes_for_messages_and_,
     test_d429de0e5338_adopt_peer_paradigm,
+    test_ec8f94139b02_codify_workspace_name_and_message_id_in_,
 )
 
 __all__ = [
@@ -31,7 +33,9 @@ __all__ = [
     "test_917195d9b5e9_add_messageembedding_table",
     "test_a1b2c3d4e5f6_initial_schema",
     "test_b765d82110bd_change_metamessages_to_user_level_with_",
+    "test_b8183c5ffb48_codify_document_level_and_times_derived",
     "test_bb6fb3a7a643_add_message_seq_in_session_column",
     "test_c3828084f472_add_indexes_for_messages_and_",
     "test_d429de0e5338_adopt_peer_paradigm",
+    "test_ec8f94139b02_codify_workspace_name_and_message_id_in_",
 ]

tests/alembic/revisions/test_b8183c5ffb48_codify_document_level_and_times_derived.py

@@ -0,0 +1,315 @@
+"""Hooks for revision b8183c5ffb48 (codify_document_level_and_times_derived)."""
+
+from __future__ import annotations
+
+import json
+
+from nanoid import generate as generate_nanoid
+from sqlalchemy import text
+
+from tests.alembic.registry import register_after_upgrade, register_before_upgrade
+from tests.alembic.verifier import MigrationVerifier
+
+# Test data constants
+WORKSPACE_NAME = "test-workspace"
+OBSERVER_NAME = "test-observer"
+OBSERVED_NAME = "test-observed"
+SESSION_NAME = "test-session"
+
+
+@register_before_upgrade("b8183c5ffb48")
+def prepare_codify_document_level_and_times_derived(
+    verifier: MigrationVerifier,
+) -> None:
+    """Seed state and assertions before upgrading to b8183c5ffb48."""
+    # Verify columns don't exist yet
+    verifier.assert_column_exists("documents", "level", exists=False)
+    verifier.assert_column_exists("documents", "times_derived", exists=False)
+
+    # Verify CHECK constraint doesn't exist yet
+    verifier.assert_constraint_exists("documents", "level_valid", "check", exists=False)
+
+    conn = verifier.conn
+    schema = verifier.schema
+
+    # Create workspace
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."workspaces" ("id", "name") VALUES (:ws_id, :ws_name)'
+        ),
+        {"ws_id": generate_nanoid(), "ws_name": WORKSPACE_NAME},
+    )
+
+    # Create peers (observer and observed)
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") '
+            + "VALUES (:observer_id, :observer_name, :ws_name), "
+            + "(:observed_id, :observed_name, :ws_name)"
+        ),
+        {
+            "observer_id": generate_nanoid(),
+            "observer_name": OBSERVER_NAME,
+            "ws_name": WORKSPACE_NAME,
+            "observed_id": generate_nanoid(),
+            "observed_name": OBSERVED_NAME,
+        },
+    )
+
+    # Create session
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."sessions" ("id", "name", "workspace_name") '
+            + "VALUES (:session_id, :session_name, :ws_name)"
+        ),
+        {
+            "session_id": generate_nanoid(),
+            "session_name": SESSION_NAME,
+            "ws_name": WORKSPACE_NAME,
+        },
+    )
+
+    # Create collection (required for documents)
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."collections" ("id", "observer", "observed", "workspace_name") '
+            + "VALUES (:collection_id, :observer, :observed, :ws_name)"
+        ),
+        {
+            "collection_id": generate_nanoid(),
+            "observer": OBSERVER_NAME,
+            "observed": OBSERVED_NAME,
+            "ws_name": WORKSPACE_NAME,
+        },
+    )
+
+    # Create documents with different scenarios
+    # Scenario 1: Document with both level and times_derived in internal_metadata
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."documents" '
+            + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") '
+            + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)"
+        ),
+        {
+            "id": generate_nanoid(),
+            "content": "Document with explicit level and times_derived=3",
+            "metadata": json.dumps({"level": "explicit", "times_derived": 3}),
+            "embedding": [0.1] * 1536,
+            "observer": OBSERVER_NAME,
+            "observed": OBSERVED_NAME,
+            "ws_name": WORKSPACE_NAME,
+            "session_name": SESSION_NAME,
+        },
+    )
+
+    # Scenario 2: Document with deductive level in internal_metadata
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."documents" '
+            + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") '
+            + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)"
+        ),
+        {
+            "id": generate_nanoid(),
+            "content": "Document with deductive level and times_derived=5",
+            "metadata": json.dumps({"level": "deductive", "times_derived": 5}),
+            "embedding": [0.2] * 1536,
+            "observer": OBSERVER_NAME,
+            "observed": OBSERVED_NAME,
+            "ws_name": WORKSPACE_NAME,
+            "session_name": SESSION_NAME,
+        },
+    )
+
+    # Scenario 3: Document without level or times_derived (should get defaults)
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."documents" '
+            + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") '
+            + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)"
+        ),
+        {
+            "id": generate_nanoid(),
+            "content": "Document without level or times_derived fields",
+            "metadata": json.dumps({"other_field": "value"}),
+            "embedding": [0.3] * 1536,
+            "observer": OBSERVER_NAME,
+            "observed": OBSERVED_NAME,
+            "ws_name": WORKSPACE_NAME,
+            "session_name": SESSION_NAME,
+        },
+    )
+
+    # Scenario 4: Document with only level in metadata
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."documents" '
+            + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") '
+            + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)"
+        ),
+        {
+            "id": generate_nanoid(),
+            "content": "Document with only level field",
+            "metadata": json.dumps({"level": "explicit"}),
+            "embedding": [0.4] * 1536,
+            "observer": OBSERVER_NAME,
+            "observed": OBSERVED_NAME,
+            "ws_name": WORKSPACE_NAME,
+            "session_name": SESSION_NAME,
+        },
+    )
+
+    # Scenario 5: Document with only times_derived in metadata
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."documents" '
+            + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") '
+            + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)"
+        ),
+        {
+            "id": generate_nanoid(),
+            "content": "Document with only times_derived field",
+            "metadata": json.dumps({"times_derived": 7}),
+            "embedding": [0.5] * 1536,
+            "observer": OBSERVER_NAME,
+            "observed": OBSERVED_NAME,
+            "ws_name": WORKSPACE_NAME,
+            "session_name": SESSION_NAME,
+        },
+    )
+
+    # Verify we have exactly 5 documents
+    count = conn.execute(text(f'SELECT COUNT(*) FROM "{schema}"."documents"')).scalar()
+    assert count == 5, f"Expected 5 documents but found {count}"
+
+
+@register_after_upgrade("b8183c5ffb48")
+def verify_codify_document_level_and_times_derived(verifier: MigrationVerifier) -> None:
+    """Add assertions validating the effects of b8183c5ffb48."""
+    # Verify columns were added with correct nullability
+    verifier.assert_column_exists("documents", "level", nullable=False)
+    verifier.assert_column_exists("documents", "times_derived", nullable=False)
+
+    # Verify CHECK constraint exists
+    verifier.assert_constraint_exists("documents", "level_valid", "check")
+
+    conn = verifier.conn
+    schema = verifier.schema
+
+    # Verify all rows have non-null values after migration
+    verifier.assert_no_nulls("documents", "level")
+    verifier.assert_no_nulls("documents", "times_derived")
+
+    # Verify data transformation: level extracted from internal_metadata
+    explicit_count = conn.execute(
+        text(f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE "level" = :level'),
+        {"level": "explicit"},
+    ).scalar()
+    deductive_count = conn.execute(
+        text(f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE "level" = :level'),
+        {"level": "deductive"},
+    ).scalar()
+
+    # 3 documents should have explicit (scenarios 1, 3, 4)
+    # 1 document should have deductive (scenario 2)
+    # 1 document should have default explicit (scenario 5)
+    assert explicit_count == 4, f"Expected 4 explicit documents, got {explicit_count}"
+    assert deductive_count == 1, f"Expected 1 deductive document, got {deductive_count}"
+
+    # Verify specific times_derived values were migrated correctly
+    times_derived_3 = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."documents" '
+            + 'WHERE "times_derived" = 3 AND "content" LIKE :pattern'
+        ),
+        {"pattern": "%explicit level and times_derived=3%"},
+    ).scalar()
+    assert (
+        times_derived_3 == 1
+    ), f"Expected 1 document with times_derived=3, got {times_derived_3}"
+
+    times_derived_5 = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."documents" '
+            + 'WHERE "times_derived" = 5 AND "content" LIKE :pattern'
+        ),
+        {"pattern": "%deductive level and times_derived=5%"},
+    ).scalar()
+    assert (
+        times_derived_5 == 1
+    ), f"Expected 1 document with times_derived=5, got {times_derived_5}"
+
+    times_derived_7 = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."documents" '
+            + 'WHERE "times_derived" = 7 AND "content" LIKE :pattern'
+        ),
+        {"pattern": "%only times_derived field%"},
+    ).scalar()
+    assert (
+        times_derived_7 == 1
+    ), f"Expected 1 document with times_derived=7, got {times_derived_7}"
+
+    # Verify default times_derived=1 was applied to documents without it
+    times_derived_1 = conn.execute(
+        text(f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE "times_derived" = 1')
+    ).scalar()
+    assert (
+        times_derived_1 == 2
+    ), f"Expected 2 documents with times_derived=1, got {times_derived_1}"
+
+    # Verify internal_metadata still contains the original data (NOT removed by migration)
+    level_in_metadata = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE internal_metadata ? \'level\''
+        )
+    ).scalar()
+    # 3 documents had level in metadata (scenarios 1, 2, 4)
+    assert (
+        level_in_metadata == 3
+    ), f"Expected 3 documents with level in metadata, got {level_in_metadata}"
+
+    times_derived_in_metadata = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE internal_metadata ? \'times_derived\''
+        )
+    ).scalar()
+    # 3 documents had times_derived in metadata (scenarios 1, 2, 5)
+    assert (
+        times_derived_in_metadata == 3
+    ), f"Expected 3 documents with times_derived in metadata, got {times_derived_in_metadata}"
+
+    # Verify server defaults work for new documents
+    new_doc_id = generate_nanoid()
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."documents" '
+            + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") '
+            + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)"
+        ),
+        {
+            "id": new_doc_id,
+            "content": "New document after migration",
+            "metadata": json.dumps({}),
+            "embedding": [0.6] * 1536,
+            "observer": OBSERVER_NAME,
+            "observed": OBSERVED_NAME,
+            "ws_name": WORKSPACE_NAME,
+            "session_name": SESSION_NAME,
+        },
+    )
+
+    # Verify the new document got default values
+    new_doc = conn.execute(
+        text(
+            f'SELECT "level", "times_derived" FROM "{schema}"."documents" WHERE "id" = :id'
+        ),
+        {"id": new_doc_id},
+    ).one()
+    assert (
+        new_doc.level == "explicit"
+    ), f"Expected new document to have level='explicit', got {new_doc.level}"
+    assert (
+        new_doc.times_derived == 1
+    ), f"Expected new document to have times_derived=1, got {new_doc.times_derived}"

tests/alembic/revisions/test_ec8f94139b02_codify_workspace_name_and_message_id_in_.py

@@ -0,0 +1,336 @@
+"""Hooks for revision ec8f94139b02 (codify_workspace_name_and_message_id_in_)."""
+
+from __future__ import annotations
+
+import json
+
+from nanoid import generate as generate_nanoid
+from sqlalchemy import text
+
+from tests.alembic.registry import register_after_upgrade, register_before_upgrade
+from tests.alembic.verifier import MigrationVerifier
+
+# Test data constants
+WORKSPACE_NAME_1 = "test-workspace-1"
+WORKSPACE_NAME_2 = "test-workspace-2"
+PEER_NAME_1 = "test-peer-1"
+PEER_NAME_2 = "test-peer-2"
+SESSION_NAME_1 = "test-session-1"
+SESSION_NAME_2 = "test-session-2"
+
+# Indexes that should be created by the migration
+_INDEXES = (
+    ("queue", "ix_queue_workspace_name"),
+    ("queue", "ix_queue_message_id_not_null"),
+    ("queue", "ix_queue_workspace_name_processed"),
+    ("queue", "ix_queue_work_unit_key_processed_id"),
+)
+
+
+@register_before_upgrade("ec8f94139b02")
+def prepare_codify_workspace_name_and_message_id_in(
+    verifier: MigrationVerifier,
+) -> None:
+    """Seed state and assertions before upgrading to ec8f94139b02."""
+    # Verify columns don't exist yet
+    verifier.assert_column_exists("queue", "workspace_name", exists=False)
+    verifier.assert_column_exists("queue", "message_id", exists=False)
+
+    # Verify foreign key constraints don't exist yet
+    verifier.assert_constraint_exists(
+        "queue", "fk_queue_workspace_name", "foreign_key", exists=False
+    )
+    verifier.assert_constraint_exists(
+        "queue", "fk_queue_message_id", "foreign_key", exists=False
+    )
+
+    # Verify indexes don't exist yet
+    verifier.assert_indexes_not_exist(_INDEXES)
+
+    conn = verifier.conn
+    schema = verifier.schema
+
+    # Create workspaces
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."workspaces" ("id", "name") '
+            + "VALUES (:ws_id_1, :ws_name_1), (:ws_id_2, :ws_name_2)"
+        ),
+        {
+            "ws_id_1": generate_nanoid(),
+            "ws_name_1": WORKSPACE_NAME_1,
+            "ws_id_2": generate_nanoid(),
+            "ws_name_2": WORKSPACE_NAME_2,
+        },
+    )
+
+    # Create peers
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") '
+            + "VALUES (:peer_id_1, :peer_name_1, :ws_name_1), "
+            + "(:peer_id_2, :peer_name_2, :ws_name_2)"
+        ),
+        {
+            "peer_id_1": generate_nanoid(),
+            "peer_name_1": PEER_NAME_1,
+            "ws_name_1": WORKSPACE_NAME_1,
+            "peer_id_2": generate_nanoid(),
+            "peer_name_2": PEER_NAME_2,
+            "ws_name_2": WORKSPACE_NAME_2,
+        },
+    )
+
+    # Create sessions
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."sessions" ("id", "name", "workspace_name") '
+            + "VALUES (:session_id_1, :session_name_1, :ws_name_1), "
+            + "(:session_id_2, :session_name_2, :ws_name_2)"
+        ),
+        {
+            "session_id_1": generate_nanoid(),
+            "session_name_1": SESSION_NAME_1,
+            "ws_name_1": WORKSPACE_NAME_1,
+            "session_id_2": generate_nanoid(),
+            "session_name_2": SESSION_NAME_2,
+            "ws_name_2": WORKSPACE_NAME_2,
+        },
+    )
+
+    # Create some messages to reference in queue items
+    message_ids: list[str] = []
+    for i in range(10):
+        message_id: str = generate_nanoid()
+        message_ids.append(message_id)
+        workspace = WORKSPACE_NAME_1 if i % 2 == 0 else WORKSPACE_NAME_2
+        session = SESSION_NAME_1 if i % 2 == 0 else SESSION_NAME_2
+        peer = PEER_NAME_1 if i % 2 == 0 else PEER_NAME_2
+
+        conn.execute(
+            text(
+                f'INSERT INTO "{schema}"."messages" '
+                + '("public_id", "workspace_name", "session_name", "peer_name", "content", "seq_in_session") '
+                + "VALUES (:msg_id, :ws_name, :session_name, :peer_name, :content, :seq)"
+            ),
+            {
+                "msg_id": message_id,
+                "ws_name": workspace,
+                "session_name": session,
+                "peer_name": peer,
+                "content": f"test message {i}",
+                "seq": i,
+            },
+        )
+
+    # Get internal message IDs for queue references
+    message_db_ids: list[int] = []
+    for msg_id in message_ids:
+        result = conn.execute(
+            text(f'SELECT "id" FROM "{schema}"."messages" WHERE "public_id" = :msg_id'),
+            {"msg_id": msg_id},
+        ).one()
+        message_db_ids.append(result.id)
+
+    # Bulk insert 100k queue items with workspace_name and message_id in payload
+    # Use efficient batch insert with generate_series
+    conn.execute(text("SET LOCAL synchronous_commit = OFF"))
+
+    # Insert queue items in three categories:
+    # 1. Items with both workspace_name and message_id (60k)
+    # 2. Items with workspace_name but NO message_id (30k)
+    # 3. Items with workspace_name and NULL message_id value (10k)
+
+    # Category 1: Both workspace_name and message_id (60k items)
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."queue" '
+            + '("work_unit_key", "task_type", "payload", "processed") '
+            + "SELECT "
+            + "  'work-unit-' || gs::text, "
+            + "  'representation', "
+            + "  jsonb_build_object("
+            + "    'workspace_name', CASE WHEN gs % 2 = 0 THEN :ws_name_1 ELSE :ws_name_2 END, "
+            + "    'message_id', :msg_db_id_0 + (gs % 10), "
+            + "    'other_field', 'value-' || gs::text"
+            + "  ), "
+            + "  false "
+            + "FROM generate_series(1, :n) AS gs"
+        ),
+        {
+            "ws_name_1": WORKSPACE_NAME_1,
+            "ws_name_2": WORKSPACE_NAME_2,
+            "msg_db_id_0": message_db_ids[0],
+            "n": 60_000,
+        },
+    )
+
+    # Category 2: Only workspace_name, no message_id key (30k items)
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."queue" '
+            + '("work_unit_key", "task_type", "payload", "processed") '
+            + "SELECT "
+            + "  'work-unit-no-msg-' || gs::text, "
+            + "  'summary', "
+            + "  jsonb_build_object("
+            + "    'workspace_name', CASE WHEN gs % 2 = 0 THEN :ws_name_1 ELSE :ws_name_2 END, "
+            + "    'other_field', 'value-' || gs::text"
+            + "  ), "
+            + "  false "
+            + "FROM generate_series(60001, :n) AS gs"
+        ),
+        {
+            "ws_name_1": WORKSPACE_NAME_1,
+            "ws_name_2": WORKSPACE_NAME_2,
+            "n": 90_000,
+        },
+    )
+
+    # Category 3: workspace_name with explicit NULL message_id (10k items)
+    conn.execute(
+        text(
+            f'INSERT INTO "{schema}"."queue" '
+            + '("work_unit_key", "task_type", "payload", "processed") '
+            + "SELECT "
+            + "  'work-unit-null-msg-' || gs::text, "
+            + "  'representation', "
+            + "  jsonb_build_object("
+            + "    'workspace_name', CASE WHEN gs % 2 = 0 THEN :ws_name_1 ELSE :ws_name_2 END, "
+            + "    'message_id', NULL::bigint, "
+            + "    'other_field', 'value-' || gs::text"
+            + "  ), "
+            + "  false "
+            + "FROM generate_series(90001, :n) AS gs"
+        ),
+        {
+            "ws_name_1": WORKSPACE_NAME_1,
+            "ws_name_2": WORKSPACE_NAME_2,
+            "n": 100_000,
+        },
+    )
+
+    # Verify we have exactly 100k queue items
+    count = conn.execute(text(f'SELECT COUNT(*) FROM "{schema}"."queue"')).scalar()
+    assert count == 100_000, f"Expected 100k queue items but found {count}"
+
+
+@register_after_upgrade("ec8f94139b02")
+def verify_codify_workspace_name_and_message_id_in(verifier: MigrationVerifier) -> None:
+    """Add assertions validating the effects of ec8f94139b02."""
+    # Verify columns were added with correct nullability
+    verifier.assert_column_exists("queue", "workspace_name", nullable=False)
+    verifier.assert_column_exists("queue", "message_id", nullable=True)
+
+    # Verify all indexes were created
+    verifier.assert_indexes_exist(_INDEXES)
+
+    # Verify foreign key constraints exist
+    verifier.assert_constraint_exists("queue", "fk_queue_workspace_name", "foreign_key")
+    verifier.assert_constraint_exists("queue", "fk_queue_message_id", "foreign_key")
+
+    conn = verifier.conn
+    schema = verifier.schema
+
+    # Verify all rows have non-null workspace_name after migration
+    verifier.assert_no_nulls("queue", "workspace_name")
+
+    # Verify data transformation: workspace_name extracted from payload
+    ws1_count = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."queue" '
+            + 'WHERE "workspace_name" = :ws_name'
+        ),
+        {"ws_name": WORKSPACE_NAME_1},
+    ).scalar()
+    ws2_count = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."queue" '
+            + 'WHERE "workspace_name" = :ws_name'
+        ),
+        {"ws_name": WORKSPACE_NAME_2},
+    ).scalar()
+
+    # Should be roughly 50/50 split (we alternate in the insert)
+    assert (
+        ws1_count == 50_000
+    ), f"Expected 50k items with workspace_name_1, got {ws1_count}"
+    assert (
+        ws2_count == 50_000
+    ), f"Expected 50k items with workspace_name_2, got {ws2_count}"
+
+    # Verify data transformation: message_id extracted from payload where it exists
+    msg_id_count = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."queue" '
+            + 'WHERE "message_id" IS NOT NULL'
+        )
+    ).scalar()
+
+    # Should be 60k items with message_id (category 1 only)
+    assert (
+        msg_id_count == 60_000
+    ), f"Expected 60k items with message_id, got {msg_id_count}"
+
+    # Verify items without message_id in payload have NULL in column
+    null_msg_id_count = conn.execute(
+        text(f'SELECT COUNT(*) FROM "{schema}"."queue" ' + 'WHERE "message_id" IS NULL')
+    ).scalar()
+
+    # Should be 40k items (30k without key + 10k with NULL value)
+    assert (
+        null_msg_id_count == 40_000
+    ), f"Expected 40k items with NULL message_id, got {null_msg_id_count}"
+
+    # Verify workspace_name was removed from payload
+    ws_in_payload_count = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."queue" '
+            + "WHERE payload ? 'workspace_name'"
+        )
+    ).scalar()
+    assert (
+        ws_in_payload_count == 0
+    ), f"Found {ws_in_payload_count} items still with workspace_name in payload"
+
+    # Verify message_id was removed from payload
+    msg_in_payload_count = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."queue" ' + "WHERE payload ? 'message_id'"
+        )
+    ).scalar()
+    assert (
+        msg_in_payload_count == 0
+    ), f"Found {msg_in_payload_count} items still with message_id in payload"
+
+    # Verify other_field remains in payload (data preservation)
+    other_field_count = conn.execute(
+        text(
+            f'SELECT COUNT(*) FROM "{schema}"."queue" '
+            + "WHERE payload ? 'other_field'"
+        )
+    ).scalar()
+    assert (
+        other_field_count == 100_000
+    ), f"Expected all 100k items to retain other_field in payload, got {other_field_count}"
+
+    # Spot check: verify a specific queue item was transformed correctly
+    sample_item = conn.execute(
+        text(
+            'SELECT "workspace_name", "message_id", "payload" '
+            + f'FROM "{schema}"."queue" '
+            + "WHERE work_unit_key = 'work-unit-1' "
+            + "LIMIT 1"
+        )
+    ).one()
+
+    assert sample_item.workspace_name == WORKSPACE_NAME_2  # gs=1 is odd, so workspace 2
+    assert sample_item.message_id is not None  # Category 1 item
+    payload = (
+        json.loads(sample_item.payload)
+        if isinstance(sample_item.payload, str)
+        else sample_item.payload
+    )
+    assert "workspace_name" not in payload
+    assert "message_id" not in payload
+    assert payload.get("other_field") == "value-1"

tests/crud/test_workspace.py

@@ -318,6 +318,7 @@ class TestWorkspaceCRUD:
             work_unit_key=f"representation:{test_workspace.name}:{session.name}:{test_peer.name}:{test_peer.name}",
             task_type="representation",
             payload={"test": "data"},
+            workspace_name=test_workspace.name,
         )
         db_session.add(queue_item)
         await db_session.flush()

tests/deriver/conftest.py

@@ -1,7 +1,7 @@
 import asyncio
-from collections.abc import Callable, Generator
+from collections.abc import Awaitable, Callable, Generator, Sequence
 from datetime import datetime, timezone
-from typing import Any, Literal
+from typing import Any, Literal, TypeAlias, cast
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -11,7 +11,10 @@ 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 get_work_unit_key
+from src.utils.work_unit import construct_work_unit_key
+
+QueuePayload: TypeAlias = dict[str, Any]
+QueuePayloadEntry: TypeAlias = QueuePayload | tuple[QueuePayload, int | None]
 
 
 @pytest.fixture
@@ -155,18 +158,29 @@ def create_queue_payload() -> Callable[..., Any]:
 @pytest.fixture
 async def add_queue_items(
     db_session: AsyncSession,
-) -> Callable[[list[dict[str, Any]], str], Any]:
+) -> Callable[
+    [Sequence[QueuePayloadEntry], str, str], Awaitable[list[models.QueueItem]]
+]:
     """Helper function to add queue items to the database"""
 
     async def _add_items(
-        payloads: list[dict[str, Any]], session_id: str
+        payloads: Sequence[QueuePayloadEntry],
+        session_id: str,
+        workspace_name: str,
     ) -> list[models.QueueItem]:
         """Add queue items to the database and return them"""
         queue_items: list[models.QueueItem] = []
-        for payload in payloads:
+        for payload_entry in payloads:
+            payload: QueuePayload
+            message_id: int | None
+            if isinstance(payload_entry, tuple):
+                payload, message_id = payload_entry
+            else:
+                payload = payload_entry
+                message_id = cast(int | None, payload.get("message_id"))
             # Generate work_unit_key from the payload
-            task_type = payload.get("task_type", "unknown")
-            work_unit_key = get_work_unit_key(payload)
+            task_type = cast(str, payload.get("task_type", "unknown"))
+            work_unit_key = construct_work_unit_key(workspace_name, payload)
 
             queue_item = models.QueueItem(
                 session_id=session_id,
@@ -174,6 +188,8 @@ async def add_queue_items(
                 work_unit_key=work_unit_key,
                 payload=payload,
                 processed=False,
+                workspace_name=workspace_name,
+                message_id=message_id,
             )
             db_session.add(queue_item)
             queue_items.append(queue_item)
@@ -203,7 +219,7 @@ async def sample_queue_items(
     messages = sample_messages
 
     # Create various types of queue payloads
-    payloads: list[dict[str, Any]] = []
+    payloads: list[tuple[dict[str, Any], int]] = []
 
     # Create representation payloads for each message
     for message in messages:
@@ -214,7 +230,7 @@ async def sample_queue_items(
             observer=message.peer_name,
             observed=message.peer_name,
         )
-        payloads.append(payload1)
+        payloads.append((payload1, message.id))
 
         # Representation for observer peer
         payload2 = create_queue_payload(
@@ -223,7 +239,7 @@ async def sample_queue_items(
             observer=peer2.name,  # peer2 observes others
             observed=message.peer_name,
         )
-        payloads.append(payload2)
+        payloads.append((payload2, message.id))
 
     # Create summary payloads for session
     for i, message in enumerate(messages):
@@ -232,10 +248,10 @@ async def sample_queue_items(
             task_type="summary",
             message_seq_in_session=i + 1,
         )
-        payloads.append(payload)
+        payloads.append((payload, message.id))
 
     # Add all payloads as queue items
-    queue_items = await add_queue_items(payloads, session.id)
+    queue_items = await add_queue_items(payloads, session.id, session.workspace_name)
 
     return queue_items
 

tests/deriver/test_deriver_processing.py

@@ -9,6 +9,7 @@ import pytest
 from src import models
 from src.deriver.deriver import process_representation_tasks_batch
 from src.utils.representation import Representation
+from src.utils.work_unit import construct_work_unit_key
 
 
 @pytest.mark.asyncio
@@ -32,14 +33,12 @@ class TestDeriverProcessing:
         sample_session_with_peers: tuple[models.Session, list[models.Peer]],
     ):
         """Test that work unit keys are generated correctly"""
-        from src.utils.work_unit import get_work_unit_key
 
         session, peers = sample_session_with_peers
         peer1, peer2, _ = peers
 
         # Create a payload for representation task
         representation_payload = {
-            "workspace_name": "workspace1",
             "session_name": session.name,
             "observer": peer2.name,
             "observed": peer1.name,
@@ -47,22 +46,25 @@ class TestDeriverProcessing:
         }
 
         # Generate work unit key for representation
-        work_unit_key = get_work_unit_key(representation_payload)
-        expected_key = (
-            f"representation:workspace1:{session.name}:{peer2.name}:{peer1.name}"
+        work_unit_key = construct_work_unit_key(
+            session.workspace_name, representation_payload
         )
+        expected_key = f"representation:{session.workspace_name}:{session.name}:{peer2.name}:{peer1.name}"
         assert work_unit_key == expected_key
 
         # Create a payload for summary task
         summary_payload = {
-            "workspace_name": "workspace1",
             "session_name": session.name,
             "task_type": "summary",
         }
 
         # Generate work unit key for summary
-        summary_work_unit_key = get_work_unit_key(summary_payload)
-        expected_summary_key = f"summary:workspace1:{session.name}:None:None"
+        summary_work_unit_key = construct_work_unit_key(
+            session.workspace_name, summary_payload
+        )
+        expected_summary_key = (
+            f"summary:{session.workspace_name}:{session.name}:None:None"
+        )
         assert summary_work_unit_key == expected_summary_key
 
     async def test_mock_queue_manager(

tests/deriver/test_queue_processing.py

@@ -10,7 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from src import models
 from src.config import settings
 from src.deriver.queue_manager import QueueManager, WorkerOwnership
-from src.utils.work_unit import get_work_unit_key
+from src.utils.work_unit import construct_work_unit_key
 
 
 @pytest.mark.asyncio
@@ -139,7 +139,7 @@ class TestQueueProcessing:
         for message in messages:
             await db_session.refresh(message)
 
-        payloads: list[Any] = []
+        payloads: list[tuple[dict[str, Any], int]] = []
         for message in messages:
             payload = create_queue_payload(  # type: ignore[reportUnknownArgumentType]
                 message=message,
@@ -147,9 +147,9 @@ class TestQueueProcessing:
                 observed=peer.name,
                 observer=peer.name,
             )
-            payloads.append(payload)
+            payloads.append((payload, message.id))
 
-        items = await add_queue_items(payloads, session.id)
+        items = await add_queue_items(payloads, session.id, session.workspace_name)
         # Determine ascending order by DB id
         ordered = (
             (
@@ -304,20 +304,23 @@ class TestQueueProcessing:
             await db_session.refresh(message)
 
         # Create queue items with token counts
-        payloads = [
-            create_queue_payload(  # type: ignore[reportUnknownArgumentType]
-                message=msg,
-                task_type="representation",
-                observed=peer.name,
-                observer=peer.name,
+        payload_entries = [
+            (
+                create_queue_payload(  # type: ignore[reportUnknownArgumentType]
+                    message=msg,
+                    task_type="representation",
+                    observed=peer.name,
+                    observer=peer.name,
+                ),
+                msg,
             )
             for msg in messages
         ]
 
         queue_items: list[models.QueueItem] = []
-        for payload in payloads:
+        for payload, message in payload_entries:
             task_type = payload.get("task_type", "unknown")
-            work_unit_key = get_work_unit_key(payload)
+            work_unit_key = construct_work_unit_key(session.workspace_name, payload)
 
             queue_item = models.QueueItem(
                 session_id=session.id,
@@ -325,6 +328,8 @@ class TestQueueProcessing:
                 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)
@@ -430,7 +435,7 @@ class TestQueueProcessing:
                 observed=peer.name,
                 observer=target.name,
             )
-            work_unit_key = get_work_unit_key(payload)
+            work_unit_key = construct_work_unit_key(session.workspace_name, payload)
 
             queue_item = models.QueueItem(
                 session_id=session.id,
@@ -438,6 +443,8 @@ class TestQueueProcessing:
                 work_unit_key=work_unit_key,
                 payload=payload,
                 processed=False,
+                workspace_name=session.workspace_name,
+                message_id=message.id,
             )
             db_session.add(queue_item)
 
@@ -596,7 +603,7 @@ class TestQueueProcessing:
                 observed=peer.name,
                 observer=target.name,
             )
-            work_unit_key = get_work_unit_key(payload)
+            work_unit_key = construct_work_unit_key(session.workspace_name, payload)
 
             queue_item = models.QueueItem(
                 session_id=session.id,
@@ -604,6 +611,8 @@ class TestQueueProcessing:
                 work_unit_key=work_unit_key,
                 payload=payload,
                 processed=False,
+                workspace_name=session.workspace_name,
+                message_id=message.id,
             )
             db_session.add(queue_item)
 
@@ -702,7 +711,6 @@ class TestQueueProcessing:
         token_counts = [500, 600]
         messages = [
             models.Message(
-                id=999,
                 session_name=session.name,
                 workspace_name=session.workspace_name,
                 peer_name=peer.name,
@@ -711,7 +719,6 @@ class TestQueueProcessing:
                 seq_in_session=1,
             ),
             models.Message(
-                id=1000,
                 session_name=session.name,
                 workspace_name=session.workspace_name,
                 peer_name=peer.name,
@@ -721,6 +728,15 @@ class TestQueueProcessing:
             ),
         ]
 
+        # Save messages to database first
+        for message in messages:
+            db_session.add(message)
+        await db_session.commit()
+
+        # Refresh to get the actual IDs
+        for message in messages:
+            await db_session.refresh(message)
+
         # Create payloads and queue items
         queue_items: list[models.QueueItem] = []
         for i, message in enumerate(messages):
@@ -729,7 +745,7 @@ class TestQueueProcessing:
             )
             payload["token_count"] = token_counts[i]
 
-            work_unit_key = get_work_unit_key(payload)
+            work_unit_key = construct_work_unit_key(session.workspace_name, payload)
 
             queue_item = models.QueueItem(
                 session_id=session.id,
@@ -737,6 +753,8 @@ class TestQueueProcessing:
                 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)
@@ -747,10 +765,11 @@ class TestQueueProcessing:
         processed_batches: list[dict[str, Any]] = []
 
         async def mock_process_item(
-            task_type: str,
-            queue_payload: dict[str, Any],  # pyright: ignore[reportUnusedParameter]
+            queue_item: models.QueueItem,
         ) -> None:
-            processed_batches.append({"task_type": task_type, "payload_count": 1})
+            processed_batches.append(
+                {"task_type": queue_item.task_type, "payload_count": 1}
+            )
 
         qm = QueueManager()
         work_unit_key = queue_items[0].work_unit_key
@@ -838,21 +857,24 @@ class TestQueueProcessing:
             await db_session.refresh(message)
 
         # Create queue items
-        payloads = [
-            create_queue_payload(  # type: ignore[reportUnknownArgumentType]
-                message=msg,
-                task_type="representation",
-                observed=peer.name,
-                observer=peer.name,
+        payload_entries = [
+            (
+                create_queue_payload(  # type: ignore[reportUnknownArgumentType]
+                    message=msg,
+                    task_type="representation",
+                    observed=peer.name,
+                    observer=peer.name,
+                ),
+                msg,
             )
             for msg in messages
         ]
 
         # Add items to queue
         queue_items: list[models.QueueItem] = []
-        for payload in payloads:
+        for payload, message in payload_entries:
             task_type = payload.get("task_type", "unknown")
-            work_unit_key = get_work_unit_key(payload)
+            work_unit_key = construct_work_unit_key(session.workspace_name, payload)
 
             queue_item = models.QueueItem(
                 session_id=session.id,
@@ -860,6 +882,8 @@ class TestQueueProcessing:
                 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)
@@ -949,21 +973,24 @@ class TestQueueProcessing:
             await db_session.refresh(message)
 
         # Create queue items
-        payloads = [
-            create_queue_payload(  # type: ignore[reportUnknownArgumentType]
-                message=msg,
-                task_type="representation",
-                observed=peer.name,
-                observer=peer.name,
+        payload_entries = [
+            (
+                create_queue_payload(  # type: ignore[reportUnknownArgumentType]
+                    message=msg,
+                    task_type="representation",
+                    observed=peer.name,
+                    observer=peer.name,
+                ),
+                msg,
             )
             for msg in messages
         ]
 
         # Add items to queue
         queue_items: list[models.QueueItem] = []
-        for payload in payloads:
+        for payload, message in payload_entries:
             task_type = payload.get("task_type", "unknown")
-            work_unit_key = get_work_unit_key(payload)
+            work_unit_key = construct_work_unit_key(session.workspace_name, payload)
 
             queue_item = models.QueueItem(
                 session_id=session.id,
@@ -971,6 +998,8 @@ class TestQueueProcessing:
                 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)

tests/integration/test_representation.py

@@ -222,8 +222,8 @@ class TestDocumentCreationWorkflow:
             observed=observed_peer.name,
             content="User likes dogs",
             session_name="test_session",
+            level="explicit",
             internal_metadata={
-                "level": "explicit",
                 "message_ids": [(1, 1)],
                 "session_name": "test_session",
             },
@@ -274,7 +274,9 @@ class TestDocumentCreationWorkflow:
             observed=observed_peer.name,
             content="Highly derived observation",
             session_name=session.name,
-            internal_metadata={"level": "explicit", "times_derived": 5},
+            level="explicit",
+            times_derived=5,
+            internal_metadata={},
             embedding=[0.1] * 1536,
         )
         db_session.add(highly_derived_doc)
@@ -286,7 +288,9 @@ class TestDocumentCreationWorkflow:
             observed=observed_peer.name,
             content="Less derived observation",
             session_name=session.name,
-            internal_metadata={"level": "explicit", "times_derived": 2},
+            level="explicit",
+            times_derived=2,
+            internal_metadata={},
             embedding=[0.2] * 1536,
         )
         db_session.add(less_derived_doc)
@@ -315,8 +319,8 @@ class TestDocumentCreationWorkflow:
             observer="test_peer",
             observed="test_peer",
             content="User said they like programming",
+            level="explicit",
             internal_metadata={
-                "level": "explicit",
                 "message_ids": [(1, 1)],
             },
             session_name="test_session",
@@ -329,8 +333,8 @@ class TestDocumentCreationWorkflow:
             observer="test_peer",
             observed="test_peer",
             content="User is likely a software developer",
+            level="deductive",
             internal_metadata={
-                "level": "deductive",
                 "message_ids": [(1, 1)],
                 "premises": ["User said they like programming"],
             },

tests/routes/test_queue_status.py

@@ -3,7 +3,7 @@ from fastapi.testclient import TestClient
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src import models
-from src.utils.work_unit import get_work_unit_key
+from src.utils.work_unit import construct_work_unit_key
 
 
 @pytest.mark.asyncio
@@ -157,9 +157,10 @@ class TestDeriverStatusEndpoint:
             queue_item = models.QueueItem(
                 session_id=session.id,
                 task_type="representation",
-                work_unit_key=get_work_unit_key(payload),
+                work_unit_key=construct_work_unit_key(workspace.name, payload),
                 payload=payload,
                 processed=False,
+                workspace_name=workspace.name,
             )
             queue_items.append(queue_item)
         db_session.add_all(queue_items)
@@ -231,9 +232,10 @@ class TestDeriverStatusEndpoint:
                 queue_item = models.QueueItem(
                     session_id=session.id,
                     task_type="representation",
-                    work_unit_key=get_work_unit_key(payload),
+                    work_unit_key=construct_work_unit_key(workspace.name, payload),
                     payload=payload,
                     processed=False,
+                    workspace_name=workspace.name,
                 )
                 queue_items.append(queue_item)
             db_session.add_all(queue_items)
@@ -292,9 +294,10 @@ class TestDeriverStatusEndpoint:
         queue_item = models.QueueItem(
             session_id=session.id,
             task_type="representation",
-            work_unit_key=get_work_unit_key(payload),
+            work_unit_key=construct_work_unit_key(workspace.name, payload),
             payload=payload,
             processed=False,
+            workspace_name=workspace.name,
         )
         db_session.add(queue_item)
         await db_session.commit()

tests/test_schema_validations.py

@@ -95,17 +95,18 @@ class TestDocumentValidations:
     def test_valid_document_create(self):
         metadata = DocumentMetadata(
             message_ids=[(1, 1)],
-            level="explicit",
             premises=[],
             message_created_at="2021-01-01T00:00:00Z",
         )
         doc = DocumentCreate(
             content="test content",
             session_name="test",
+            level="explicit",
             metadata=metadata,
             embedding=[0.1, 0.2, 0.3],
         )
         assert doc.content == "test content"
+        assert doc.level == "explicit"
         assert doc.metadata == metadata
 
     def test_document_content_too_short(self):
@@ -113,9 +114,9 @@ class TestDocumentValidations:
             DocumentCreate(
                 content="",
                 session_name="test",
+                level="explicit",
                 metadata=DocumentMetadata(
                     message_ids=[(1, 1)],
-                    level="explicit",
                     premises=[],
                     message_created_at="2021-01-01T00:00:00Z",
                 ),
@@ -129,9 +130,9 @@ class TestDocumentValidations:
             DocumentCreate(
                 content="a" * 100001,
                 session_name="test",
+                level="explicit",
                 metadata=DocumentMetadata(
                     message_ids=[(1, 1)],
-                    level="explicit",
                     premises=[],
                     message_created_at="2021-01-01T00:00:00Z",
                 ),