Bug: extract_thread_messages can pull in an unrelated regeneration-root message and truncate the real thread early
- Dominant language
- TypeScript
- Stars
- 156k
- Forks
- 24.6k
- Avg merge
- 20h 50m
- Merged PRs (30d)
- 586
Description
### Self Checks
- [x] I have read the Contributing Guide and Language Policy.
- [x] This is only for bug report; questions go to Discussions.
- [x] I have searched for existing issues, including closed ones.
- [x] I confirm that I am using English to submit this report.
### Dify version
main (latest)
### Cloud or Self Hosted
Self-hosted / source-level bug in API.
### Steps to reproduce
`api/core/prompt/utils/extract_thread_messages.py`:
```python
def extract_thread_messages(messages: Sequence[Message]):
thread_messages: list[Message] = []
next_message = None
for message in messages:
if not message.parent_message_id:
# If the message is regenerated and does not have a parent message, it is the start of a new thread
thread_messages.append(message)
break
if not next_message:
thread_messages.append(message)
next_message = message.parent_message_id
else:
if next_message in {message.id, UUID_NIL}:
thread_messages.append(message)
next_message = message.parent_message_id
return thread_messages
```
The `if not message.parent_message_id:` check runs unconditionally for **every** message in the loop, not just the one currently being tracked via `next_message`. Per the comment above it, a message with no `parent_message_id` (as opposed to the legacy `UUID_NIL` sentinel) marks "the start of a new thread" — i.e. a regeneration point. A single conversation can contain **multiple** such regeneration-root messages (one per regenerate action over the conversation's life), only one of which is actually the ancestor of the thread currently being walked.
If an *unrelated* regeneration-root message (not equal to the current `next_message` target) happens to sit between the message we're walking from and its real parent in the `created_at DESC`-ordered list, this check fires on it first — the walk incorrectly grabs that unrelated message and `break`s, silently dropping every real ancestor beyond it (and including a message that never belonged to this thread).
Minimal reproduction (pure function, no DB needed — mirrors the real `created_at DESC` message order from `get_thread_messages_length`'s query):
```python
from dataclasses import dataclass
UUID_NIL = "00000000-0000-0000-0000-000000000000"
@dataclass
class Message:
id: str
parent_message_id: str | None
answer: str = "some answer"
def extract_thread_messages(messages):
thread_messages = []
next_message = None
for message in messages:
if not message.parent_message_id:
thread_messages.append(message)
break
if not next_message:
thread_messages.append(message)
next_message = message.parent_message_id
else:
if next_message in {message.id, UUID_NIL}:
thread_messages.append(message)
next_message = message.parent_message_id
return thread_messages
# Real thread: C -> B -> A (A is the true root).
# X is an unrelated, earlier regeneration-root message from the same conversation
# (parent_message_id=None), created chronologically between B and C but NOT an
# ancestor of C.
msgC = Message(id="C", parent_message_id="B")
msgX = Message(id="X", parent_message_id=None)
msgB = Message(id="B", parent_message_id="A")
msgA = Message(id="A", parent_message_id=None)
messages = [msgC, msgX, msgB, msgA] # created_at DESC order
print([m.id for m in extract_thread_messages(messages)])
```
Output: `['C', 'X']`
### ✔️ Expected Behavior
`['C', 'B', 'A']` — `X` is unrelated to the thread being walked from `C` and should be skipped, not treated as the thread's root.
### ❌ Actual Behavior
`extract_thread_messages` returns `['C', 'X']`: it wrongly appends the unrelated regeneration-root message `X` and stops the walk there, losing the real ancestors `B` and `A`.
### Impact
`extract_thread_messages` is the shared thread-reconstruction primitive used by:
- `core/memory/token_buffer_memory.py` (conversation history sent to the LLM)
- `core/agent/base_agent_runner.py` (agent conversation history)
- `core/prompt/utils/get_thread_messages_length.py` → `core/app/apps/advanced_chat/app_generator.py`'s `dialogue_count`
A corrupted/truncated thread here can silently shrink the conversation history an LLM call actually sees, and undercounts `dialogue_count`, in any conversation that has had more than one regeneration.
### Proposed fix direction
Only treat "no `parent_message_id`" as the thread's true root when the message being examined is the one actually being tracked (i.e. scope the check to the branch that already matched via `next_message`), rather than checking it unconditionally on every message in the loop. Happy to submit a PR with a regression test (mirroring the existing `test_extract_thread_messages_mixed_with_legacy_messages`-style tests in `api/tests/unit_tests/core/prompt/test_extract_thread_messages.py`) once a maintainer confirms the intended semantics — in particular, whether the `None`-vs-`UUID_NIL` distinction for "start of new thread" is meant to always take precedence, or should only apply when it's actually the message we're chasing.
---
Disclosure: this bug was found and reproduced with AI assistance (Claude Code); I verified the standalone repro above executes exactly as shown before filing.
Contributor guide
Research direction
Start with api/core/prompt/utils/extract_thread_messages.py and compare its behavior with the tests in api/tests/unit_tests/core/prompt/test_extract_thread_messages.py, especially the mixed legacy-message cases. Add a regression case based on the C/X/B/A reproduction, then run that test file and confirm unrelated regeneration roots are excluded while the real ancestors remain.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend-api-design
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 72/100