feat(db): add conversation_id to conversation_items primary key

Widen the conversation_items primary key from (workspace_id, id) to
(workspace_id, conversation_id, id) so a conversation's items stay
contiguous under the workspace prefix for the per-conversation prefix
scans that dominate item reads.

Co-authored-by: Isaac

omnigent/db/db_models.py

@@ -602,10 +602,13 @@ class SqlConversationItem(Base):
         server_default="0",
         default=current_workspace_id,
     )
-    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    # conversation_id leads id in the PK so a conversation's items stay
+    # contiguous for the per-conversation prefix scans that dominate reads.
     conversation_id: Mapped[str] = mapped_column(
         String(64),
+        primary_key=True,
     )
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
     response_id: Mapped[str] = mapped_column(String(64))
     created_at: Mapped[int] = mapped_column(Integer)
     # Enum stored as a stable int code (see omnigent.db.enum_codecs

omnigent/db/migrations/versions/y1a2b3c4d5e6_widen_conversation_items_pk_with_conversation_id.py

@@ -0,0 +1,97 @@
+"""Add conversation_id to the conversation_items primary key.
+
+Revision ID: y1a2b3c4d5e6
+Revises: x1a2b3c4d5e6
+Create Date: 2026-07-08 00:00:00.000000
+
+Widens the ``conversation_items`` primary key from ``(workspace_id, id)``
+to ``(workspace_id, conversation_id, id)``.  ``conversation_id`` slots in
+between the tenant partition key and the item id so a single conversation's
+items stay contiguous under the workspace prefix, matching the per-conversation
+prefix scans that dominate item reads.  ``conversation_id`` is already NOT NULL
+and every existing row has one, so the rebuild is a pure key change with no
+backfill.
+
+There are no FK constraints in the schema (see ``p1a2b3c4d5e6``), so rebuilding
+the primary key is a purely local operation on this one table.
+
+SQLite note: ``batch_alter_table(recreate="always")`` rebuilds the table so the
+primary key can change (SQLite cannot alter a PK in place); the new
+``create_primary_key`` overrides the reflected key.  On PostgreSQL the existing
+named PK is dropped explicitly first (a table can hold only one primary key)
+before the wider one is added.  Both paths guard the rebuild with
+``PRAGMA foreign_keys`` on SQLite.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import warnings
+from collections.abc import Iterator, Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+revision: str = "y1a2b3c4d5e6"
+down_revision: str | None = "x1a2b3c4d5e6"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+_TABLE = "conversation_items"
+# Primary key before this migration and after it.
+_OLD_PK = ["workspace_id", "id"]
+_NEW_PK = ["workspace_id", "conversation_id", "id"]
+
+
+def _is_sqlite() -> bool:
+    return op.get_bind().dialect.name == "sqlite"
+
+
+def _existing_pk_name(table: str) -> str | None:
+    """Reflect the current primary-key constraint name (PostgreSQL path)."""
+    return sa.inspect(op.get_bind()).get_pk_constraint(table).get("name")
+
+
+@contextlib.contextmanager
+def _quiet_pk_override() -> Iterator[None]:
+    """
+    Silence the expected SQLite batch-rebuild warning about the reflected
+    primary key not matching the wider one we install. The override is
+    intentional here, and this fires on every fresh DB.
+    """
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r".*not matching locally specified columns.*",
+            category=sa.exc.SAWarning,
+        )
+        yield
+
+
+def _rebuild_pk(new_pk: list[str]) -> None:
+    """Drop the current ``conversation_items`` PK and install ``new_pk``."""
+    sqlite = _is_sqlite()
+    if sqlite:
+        op.execute(sa.text("PRAGMA foreign_keys = OFF"))
+
+    old_pk_name = None if sqlite else _existing_pk_name(_TABLE)
+    with (
+        _quiet_pk_override(),
+        op.batch_alter_table(_TABLE, recreate="always" if sqlite else "auto") as batch_op,
+    ):
+        if old_pk_name is not None:
+            batch_op.drop_constraint(old_pk_name, type_="primary")
+        batch_op.create_primary_key(f"pk_{_TABLE}", new_pk)
+
+    if sqlite:
+        op.execute(sa.text("PRAGMA foreign_keys = ON"))
+
+
+def upgrade() -> None:
+    """Widen the primary key to ``(workspace_id, conversation_id, id)``."""
+    _rebuild_pk(_NEW_PK)
+
+
+def downgrade() -> None:
+    """Restore the ``(workspace_id, id)`` primary key."""
+    _rebuild_pk(_OLD_PK)

tests/db/test_db_models.py

@@ -444,7 +444,7 @@ class TestSqlConversationItem:
             session.add(item)
 
         with managed() as session:
-            loaded = session.get(SqlConversationItem, (0, "msg_test1"))
+            loaded = session.get(SqlConversationItem, (0, "conv_test1", "msg_test1"))
             assert loaded is not None
             assert loaded.conversation_id == "conv_test1"
             assert loaded.type == encode_item_type("message")
@@ -489,7 +489,7 @@ class TestSqlConversationItem:
 
         # Without FK cascade the item is NOT automatically deleted.
         with managed() as session:
-            assert session.get(SqlConversationItem, (0, "msg_del")) is not None
+            assert session.get(SqlConversationItem, (0, "conv_del", "msg_del")) is not None
 
     def test_multiple_items_ordered_by_position(self, db_uri: str) -> None:
         engine = get_or_create_engine(db_uri)

tests/db/test_migration_conversation_items_pk.py

@@ -0,0 +1,84 @@
+"""Tests for the conversation_items PK-widening migration (y1a2b3c4d5e6)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import sqlalchemy as sa
+from alembic import command
+from sqlalchemy.engine import Engine
+
+from omnigent.db.utils import (
+    _build_alembic_config,
+    clear_engine_cache,
+    get_or_create_engine,
+)
+
+_TABLE = "conversation_items"
+_HEAD_PK = ["workspace_id", "conversation_id", "id"]
+_PRIOR_PK = ["workspace_id", "id"]
+
+
+def _pk(engine: Engine) -> list[str]:
+    return sa.inspect(engine).get_pk_constraint(_TABLE)["constrained_columns"]
+
+
+def test_head_widens_conversation_items_pk(tmp_path: Path) -> None:
+    """At head the PK is ``(workspace_id, conversation_id, id)``."""
+    uri = f"sqlite:///{tmp_path / 'head.db'}"
+    engine = get_or_create_engine(uri)
+    try:
+        assert _pk(engine) == _HEAD_PK
+    finally:
+        engine.dispose()
+        clear_engine_cache()
+
+
+def test_downgrade_restores_prior_conversation_items_pk(tmp_path: Path) -> None:
+    """Downgrading one step drops conversation_id back out of the PK."""
+    uri = f"sqlite:///{tmp_path / 'downgrade.db'}"
+    engine = get_or_create_engine(uri)
+    try:
+        assert _pk(engine) == _HEAD_PK
+
+        config = _build_alembic_config(uri)
+        with engine.begin() as conn:
+            config.attributes["connection"] = conn
+            command.downgrade(config, "x1a2b3c4d5e6")
+
+        assert _pk(engine) == _PRIOR_PK
+    finally:
+        engine.dispose()
+        clear_engine_cache()
+
+
+def test_items_round_trip_via_store(tmp_path: Path) -> None:
+    """Append + list still works with conversation_id in the PK."""
+    from omnigent.entities.conversation import MessageData, NewConversationItem
+    from omnigent.stores.conversation_store.sqlalchemy_store import (
+        SqlAlchemyConversationStore,
+    )
+
+    uri = f"sqlite:///{tmp_path / 'roundtrip.db'}"
+    engine = get_or_create_engine(uri)
+    try:
+        store = SqlAlchemyConversationStore(str(engine.url))
+        conv = store.create_conversation()
+        store.append(
+            conv.id,
+            [
+                NewConversationItem(
+                    type="message",
+                    response_id="resp_1",
+                    data=MessageData(
+                        role="user",
+                        content=[{"type": "input_text", "text": "hi"}],
+                    ),
+                )
+            ],
+        )
+        items = store.list_items(conv.id).data
+        assert [i.type for i in items] == ["message"]
+    finally:
+        engine.dispose()
+        clear_engine_cache()

tests/db/test_migration_workspace_id.py

@@ -40,12 +40,15 @@ _ORIGINAL_PKS: dict[str, list[str]] = {
 
 # Tables whose PK was changed again by a migration that came after
 # r1a2b3c4d5e6. The value is the expected PK at the current head.
-# ``test_workspace_id_leads_the_primary_key`` skips these so that later
+# ``test_workspace_id_leads_the_primary_key`` uses these so that later
 # migrations don't break the r-migration test.
 _LATER_PK_OVERRIDES: dict[str, list[str]] = {
     # v1a2b3c4d5e6 replaced (workspace_id, owner, name) with
     # (workspace_id, host_id) for the hosts table.
     "hosts": ["workspace_id", "host_id"],
+    # y1a2b3c4d5e6 widened conversation_items to insert conversation_id
+    # between workspace_id and id.
+    "conversation_items": ["workspace_id", "conversation_id", "id"],
 }