langgenius / langgenius/dify

Bug: extract_thread_messages can pull in an unrelated regeneration-root message and truncate the real thread early

Open Beginner friendly
#39,031 1 comment 1 reaction 0 assignees View on GitHub
project#dify
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.