Skip to content

fix: delete steps and elements on message edit#2915

Open
Allaoua9 wants to merge 6 commits into
Chainlit:mainfrom
Allaoua9:fix/reset-steps-message-edit-al
Open

fix: delete steps and elements on message edit#2915
Allaoua9 wants to merge 6 commits into
Chainlit:mainfrom
Allaoua9:fix/reset-steps-message-edit-al

Conversation

@Allaoua9
Copy link
Copy Markdown
Contributor

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.

This is based on this PR .

cc @hayescode this PR addresses your concern on orphaned Elements and Feedbacks.

@dosubot dosubot Bot added size:M This PR changes 30-99 lines, ignoring generated files. bug Something isn't working unit-tests Has unit tests. labels Apr 30, 2026
@Allaoua9
Copy link
Copy Markdown
Contributor Author

PR 2856 should be closed in favor of this.

Copy link
Copy Markdown
Contributor

@cubic-dev-ai cubic-dev-ai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 issues found across 3 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/tests/test_message.py">

<violation number="1" location="backend/tests/test_message.py:897">
P2: This assertion is too loose and can miss extra delete_* calls, so the test may not catch over-deletion regressions.</violation>

<violation number="2" location="backend/tests/test_message.py:897">
P2: Deletion tests check async data-layer calls with call assertions instead of await assertions, so they would miss a missing-`await` regression.</violation>
</file>

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

Comment thread backend/tests/test_message.py Outdated
Comment thread backend/tests/test_message.py
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR fixes persistence-layer cleanup when a user edits a message by ensuring descendant steps (and their associated elements/feedback) from prior executions are deleted, preventing historical artifacts from accumulating and resurfacing in the UI after edits.

Changes:

  • Add Message.remove_children() to delete descendant steps leaves-first, plus associated feedback and elements.
  • Invoke remove_children() during message edit handling and message removal.
  • Add unit tests covering no data layer, missing thread, no children, nested descendants, and feedback/element cleanup ordering.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File Description
backend/tests/test_message.py Adds test coverage for descendant cleanup behavior of remove_children().
backend/chainlit/socket.py Calls remove_children() during edit_message to clear previous execution artifacts.
backend/chainlit/message.py Introduces remove_children() and invokes it from remove() to delete descendants and related persisted data.

Comment on lines 132 to +146
data_layer = get_data_layer()
if data_layer:
try:
asyncio.create_task(data_layer.delete_step(step_dict["id"]))
except Exception as e:
if self.fail_on_persist_error:
raise e
logger.error(f"Failed to persist message deletion: {e!s}")

try:
await self.remove_children()
except Exception as e:
if self.fail_on_persist_error:
raise e
logger.error(f"Failed to persist message children deletion: {e!s}")
Comment on lines +161 to +179
steps = thread.get("steps", [])

def collect_descendants(parent_id: str, visited: Optional[set] = None) -> list:
"""Return descendant IDs in post-order (leaves first, parents last)."""
if visited is None:
visited = set()
if parent_id in visited:
return []
visited.add(parent_id)
result = []
for step in steps:
if step.get("parentId") == parent_id:
result.extend(collect_descendants(step["id"], visited))
result.append(step["id"])
return result

# Ordered leaves-first so that referential integrity constraints are respected.
ordered_descendant_ids = collect_descendants(self.id)
descendant_set = set(ordered_descendant_ids)
Comment on lines 336 to 340
if message.id == payload["message"]["id"]:
message.content = payload["message"]["output"]
await message.update()
await message.remove_children()
orig_message = message
Copy link
Copy Markdown
Collaborator

@dokterbob dokterbob left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the contrib @Allaoua9! This is quite a complex piece of code and will take us a while to properly review.

Until that time, could you please look at copilot's feedback? It might well be BS but it also might not be. It often times isn't. ;)

A good way to approach this, if you haven't already, is to run the unit test with branch coverage and see whether there might be branches you've missed. Case like this, with a bit of vibe coding, 100% coverage seems feasible.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working size:M This PR changes 30-99 lines, ignoring generated files. unit-tests Has unit tests.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants