backend/alembic/versions/051_extend_ai_chat_session.py

@@ -0,0 +1,39 @@
+"""extend ai chat session with tree_id and archived_at
+
+Revision ID: 051
+Revises: 050
+"""
+from alembic import op
+import sqlalchemy as sa
+
+revision = "051"
+down_revision = "050"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "ai_chat_sessions",
+        sa.Column("tree_id", sa.UUID(), nullable=True),
+    )
+    op.add_column(
+        "ai_chat_sessions",
+        sa.Column("archived_at", sa.DateTime(timezone=True), nullable=True),
+    )
+    op.create_index("ix_ai_chat_sessions_tree_id", "ai_chat_sessions", ["tree_id"])
+    op.create_foreign_key(
+        "fk_ai_chat_sessions_tree_id",
+        "ai_chat_sessions",
+        "trees",
+        ["tree_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+
+def downgrade() -> None:
+    op.drop_constraint("fk_ai_chat_sessions_tree_id", "ai_chat_sessions", type_="foreignkey")
+    op.drop_index("ix_ai_chat_sessions_tree_id", table_name="ai_chat_sessions")
+    op.drop_column("ai_chat_sessions", "archived_at")
+    op.drop_column("ai_chat_sessions", "tree_id")

backend/app/models/ai_chat_session.py

@@ -86,3 +86,14 @@ class AIChatSession(Base):
         default=lambda: datetime.now(timezone.utc),
         onupdate=lambda: datetime.now(timezone.utc),
     )
+    # Editor-embedded session: links to a specific tree/flow
+    tree_id: Mapped[Optional[uuid.UUID]] = mapped_column(
+        UUID(as_uuid=True),
+        ForeignKey("trees.id", ondelete="CASCADE"),
+        nullable=True,
+        index=True,
+    )
+    archived_at: Mapped[Optional[datetime]] = mapped_column(
+        DateTime(timezone=True),
+        nullable=True,
+    )