fix: delete steps and elements on message edit

[tonca]

Problem

When a user edits a message, child steps and elements from previous executions were never deleted from the data layer. This caused all historical steps to accumulate and reappear in the UI after each edit.

---

Summary by cubic

Deletes descendant steps when a message is edited or removed, preventing old executions from reappearing and keeping the thread clean. Elements are not deleted.

• Bug Fixes
  • Added Message.remove_children() to recursively delete descendant steps in leaves-first order, never deleting the root.
  • Invoked remove_children() on edit (socket.edit_message) and delete (Message.remove) to keep data consistent.
  • Expanded tests to cover no data layer, missing thread, no children, nested descendants, and to confirm elements are untouched.

^{Written for commit 92b3399. Summary will update on new commits.}

[cubic-dev-ai]

3 issues found across 4 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="backend/chainlit/message.py">

<violation number="1" location="backend/chainlit/message.py:141">
P2: Unhandled exceptions in `remove_children()` can abort `remove()` before UI deletion and ignore best-effort persistence semantics when `fail_on_persist_error` is false.</violation>
</file>

<file name="backend/tests/test_message.py">

<violation number="1" location="backend/tests/test_message.py:828">
P2: Set-based delete-step assertion hides call order and duplicate deletions, weakening regression coverage for ordered descendant cleanup.</violation>

<violation number="2" location="backend/tests/test_message.py:878">
P2: AsyncMock deletion methods are asserted with call-based checks instead of await-based checks, so tests can miss regressions where deletions are no longer awaited.</violation>
</file>

---

Since this is your first cubic review, here's how it works:

• cubic automatically reviews your code and comments on bugs and improvements
• Teach cubic by replying to its comments. cubic learns from your replies and gets better over time
• Add one-off context when rerunning by tagging @cubic-dev-ai with guidance or docs links (including llms.txt)
• Ask questions if you need clarification on any suggestion

_{Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.}

[cubic-dev-ai]

1 issue found across 4 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="backend/chainlit/message.py">

<violation number="1" location="backend/chainlit/message.py:141">
P1: Unhandled exceptions in `remove_children()` can abort `remove()` and prevent UI step deletion, violating the class’s persistence-error handling behavior.</violation>
</file>

_{Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.}

[Allaoua9]

since data layers already delete the elements in delete_step method. there might be no need to call delete_element here right ? @cubic-dev-ai

[cubic-dev-ai]

@Allaoua9 You're likely correct that for most standard data layer implementations (like the LiteralDataLayer), calling delete_step already triggers a cascading deletion of associated elements. In those cases, the explicit call to delete_element is indeed redundant.

However, there are a few reasons why this explicit cleanup is likely included:

1. Interface Guarantees: The BaseDataLayer interface does not strictly enforce that delete_step must handle element cleanup. By explicitly deleting the elements in the Message class, Chainlit ensures that custom data layer implementations (e.g., a simple file-based or custom SQL implementation) don't accidentally leave "orphaned" elements behind if they haven't implemented cascading deletes.
2. Completeness: While delete_step might handle elements directly tied to that specific step, the remove_children logic is designed to be a "clean sweep" of the entire descendant tree. Including the elements explicitly makes the intent of the method clear and ensures nothing is missed regardless of the backend's internal logic.
3. UI Sync: It ensures that the elements are cleared out logically before the final emitter.delete_step call notifies the frontend, providing a consistent state.

In short: it's a safeguard for consistency across different data layer implementations. If your data layer already handles this, these extra calls will simply be no-ops or redundant 404s, but they prevent leaks for users with simpler custom backends.

One thing to note: delete_step definitely does not recursively delete child steps (descendants in the tree) in most backends, which is why the recursive collect_descendants logic for the steps themselves is absolutely necessary.

Chainlit Data Layer Documentation

[hayescode]

Please make descendant cleanup consistent across all supported data layers before approval. Right now the new deletion path can still leave orphaned persisted records behind.

[hayescode]

This is a persistence-affecting change, so we need consistent behavior across all supported data layers. Right now this path relies on delete_step(), but DynamoDBDataLayer.delete_step() only deletes the step item and does not clean up attached elements/feedback. That means edits can still leave orphaned persisted records behind. Please either delete descendant attachments explicitly here or update every affected data layer and test coverage to support the new semantics.