From 932b418681efc0714f4e88af0a94b5f3e35d0fc7 Mon Sep 17 00:00:00 2001
From: Matthew Zhou <mattzh1314@gmail.com>
Date: Tue, 7 Oct 2025 13:26:05 -0700
Subject: [PATCH] feat: Add FK cascades, remove app-side label propagation
 [LET-4690] (#5219)

add FK cascades, remove app-side label propagation
---
 ..._add_cascades_to_blocks_agents_fks_set_.py | 55 +++++++++++++++++++
 letta/orm/block.py                            | 17 +-----
 letta/orm/blocks_agents.py                    | 10 +++-
 3 files changed, 64 insertions(+), 18 deletions(-)
 create mode 100644 alembic/versions/038e68cdf0df_add_cascades_to_blocks_agents_fks_set_.py

diff --git a/alembic/versions/038e68cdf0df_add_cascades_to_blocks_agents_fks_set_.py b/alembic/versions/038e68cdf0df_add_cascades_to_blocks_agents_fks_set_.py
new file mode 100644
index 00000000..83406ef5
--- /dev/null
+++ b/alembic/versions/038e68cdf0df_add_cascades_to_blocks_agents_fks_set_.py
@@ -0,0 +1,55 @@
+"""add cascades to blocks_agents FKs; set initially immediate
+
+Revision ID: 038e68cdf0df
+Revises: b6061da886ee
+Create Date: 2025-10-07 13:01:17.872405
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "038e68cdf0df"
+down_revision: Union[str, None] = "b6061da886ee"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(op.f("blocks_agents_agent_id_fkey"), "blocks_agents", type_="foreignkey")
+    op.drop_constraint(op.f("fk_block_id_label"), "blocks_agents", type_="foreignkey")
+    op.create_foreign_key(
+        "fk_block_id_label",
+        "blocks_agents",
+        "block",
+        ["block_id", "block_label"],
+        ["id", "label"],
+        onupdate="CASCADE",
+        ondelete="CASCADE",
+        initially="IMMEDIATE",
+        deferrable=True,
+    )
+    op.create_foreign_key(None, "blocks_agents", "agents", ["agent_id"], ["id"], ondelete="CASCADE")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, "blocks_agents", type_="foreignkey")
+    op.drop_constraint("fk_block_id_label", "blocks_agents", type_="foreignkey")
+    op.create_foreign_key(
+        op.f("fk_block_id_label"),
+        "blocks_agents",
+        "block",
+        ["block_id", "block_label"],
+        ["id", "label"],
+        initially="DEFERRED",
+        deferrable=True,
+    )
+    op.create_foreign_key(op.f("blocks_agents_agent_id_fkey"), "blocks_agents", "agents", ["agent_id"], ["id"])
+    # ### end Alembic commands ###
diff --git a/letta/orm/block.py b/letta/orm/block.py
index de5eafc3..4ae5e576 100644
--- a/letta/orm/block.py
+++ b/letta/orm/block.py
@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING, List, Optional, Type
 
 from sqlalchemy import JSON, BigInteger, ForeignKey, Index, Integer, String, UniqueConstraint, event
-from sqlalchemy.orm import Mapped, attributes, declared_attr, mapped_column, relationship
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship
 
 from letta.constants import CORE_MEMORY_BLOCK_CHAR_LIMIT
 from letta.orm.block_history import BlockHistory
@@ -110,21 +110,6 @@ class Block(OrganizationMixin, SqlalchemyBase, ProjectMixin, TemplateEntityMixin
         )  # Helps manage potential FK cycles
 
 
-@event.listens_for(Block, "after_update")  # Changed from 'before_update'
-def block_before_update(mapper, connection, target):
-    """Handle updating BlocksAgents when a block's label changes."""
-    label_history = attributes.get_history(target, "label")
-    if not label_history.has_changes():
-        return
-
-    blocks_agents = BlocksAgents.__table__
-    connection.execute(
-        blocks_agents.update()
-        .where(blocks_agents.c.block_id == target.id, blocks_agents.c.block_label == label_history.deleted[0])
-        .values(block_label=label_history.added[0])
-    )
-
-
 @event.listens_for(Block, "before_insert")
 @event.listens_for(Block, "before_update")
 def validate_value_length(mapper, connection, target):
diff --git a/letta/orm/blocks_agents.py b/letta/orm/blocks_agents.py
index 918363d8..6c937428 100644
--- a/letta/orm/blocks_agents.py
+++ b/letta/orm/blocks_agents.py
@@ -15,7 +15,13 @@ class BlocksAgents(Base):
             name="unique_label_per_agent",
         ),
         ForeignKeyConstraint(
-            ["block_id", "block_label"], ["block.id", "block.label"], name="fk_block_id_label", deferrable=True, initially="DEFERRED"
+            ["block_id", "block_label"],
+            ["block.id", "block.label"],
+            name="fk_block_id_label",
+            onupdate="CASCADE",
+            ondelete="CASCADE",
+            deferrable=True,
+            initially="IMMEDIATE",
         ),
         UniqueConstraint("agent_id", "block_id", name="unique_agent_block"),
         Index("ix_blocks_agents_block_label_agent_id", "block_label", "agent_id"),
@@ -24,6 +30,6 @@ class BlocksAgents(Base):
     )
 
     # unique agent + block label
-    agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id"), primary_key=True)
+    agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id", ondelete="CASCADE"), primary_key=True)
     block_id: Mapped[str] = mapped_column(String, primary_key=True)
     block_label: Mapped[str] = mapped_column(String, primary_key=True)