mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
4e9f6c68ea
* fix(persistence): retry feedback and conversation upserts that lose a race FeedbackRepository.upsert and ChannelConnectionRepository.set_thread_id looked up the row and inserted it when missing, with no conflict handling. Two first writers for one key both missed the lookup, and the second INSERT violated uq_feedback_thread_run_user or uq_channel_conversation_connection_external: a 500 for the second of two concurrent first ratings, and a failed message when a conversation's first message raced /new. Both now catch IntegrityError, roll back and retry, so the losing writer re-reads the committed row and updates it, as upsert_connection already does. A retry rather than ON CONFLICT keeps the NULL-user feedback path updating its single row, since a NULL key never conflicts. * docs(changelog): reference #6852 for the upsert race fix * test(persistence): pin upsert retry exhaustion and share the miss-select double Both retry loops now have a test that makes every SELECT miss against a pre-seeded row, so all attempts hit the unique key: the original IntegrityError must propagate after exactly _UPSERT_MAX_ATTEMPTS attempts and the stored row must stay unchanged. The "first SELECT misses" session double lived in four copies (one per retry test, including upsert_connection's). It now lives once in tests/_repository_race_helpers.py, with a misses count so the same double drives both the single lost race and exhaustion.
415 lines
17 KiB
Python
415 lines
17 KiB
Python
"""Tests for FeedbackRepository and follow-up association.
|
|
|
|
Uses temp SQLite DB for ORM tests.
|
|
"""
|
|
|
|
import pytest
|
|
from _repository_race_helpers import miss_select_factory
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
from deerflow.persistence.feedback import FeedbackRepository
|
|
from deerflow.persistence.feedback.sql import _UPSERT_MAX_ATTEMPTS
|
|
|
|
|
|
async def _make_feedback_repo(tmp_path):
|
|
from deerflow.persistence.engine import get_session_factory, init_engine
|
|
|
|
url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}"
|
|
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
|
|
return FeedbackRepository(get_session_factory())
|
|
|
|
|
|
async def _cleanup():
|
|
from deerflow.persistence.engine import close_engine
|
|
|
|
await close_engine()
|
|
|
|
|
|
# -- FeedbackRepository --
|
|
|
|
|
|
class TestFeedbackRepository:
|
|
@pytest.mark.anyio
|
|
async def test_create_positive(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
record = await repo.create(run_id="r1", thread_id="t1", rating=1)
|
|
assert record["feedback_id"]
|
|
assert record["rating"] == 1
|
|
assert record["run_id"] == "r1"
|
|
assert record["thread_id"] == "t1"
|
|
assert "created_at" in record
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_create_negative_with_comment(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
record = await repo.create(
|
|
run_id="r1",
|
|
thread_id="t1",
|
|
rating=-1,
|
|
comment="Response was inaccurate",
|
|
)
|
|
assert record["rating"] == -1
|
|
assert record["comment"] == "Response was inaccurate"
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_create_with_message_id(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
record = await repo.create(run_id="r1", thread_id="t1", rating=1, message_id="msg-42")
|
|
assert record["message_id"] == "msg-42"
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_create_with_owner(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
record = await repo.create(run_id="r1", thread_id="t1", rating=1, user_id="user-1")
|
|
assert record["user_id"] == "user-1"
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_create_invalid_rating_zero(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
with pytest.raises(ValueError):
|
|
await repo.create(run_id="r1", thread_id="t1", rating=0)
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_create_invalid_rating_five(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
with pytest.raises(ValueError):
|
|
await repo.create(run_id="r1", thread_id="t1", rating=5)
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_get(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
created = await repo.create(run_id="r1", thread_id="t1", rating=1)
|
|
fetched = await repo.get(created["feedback_id"])
|
|
assert fetched is not None
|
|
assert fetched["feedback_id"] == created["feedback_id"]
|
|
assert fetched["rating"] == 1
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_get_nonexistent(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
assert await repo.get("nonexistent") is None
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_by_run(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
await repo.create(run_id="r1", thread_id="t1", rating=1, user_id="user-1")
|
|
await repo.create(run_id="r1", thread_id="t1", rating=-1, user_id="user-2")
|
|
await repo.create(run_id="r2", thread_id="t1", rating=1, user_id="user-1")
|
|
results = await repo.list_by_run("t1", "r1", user_id=None)
|
|
assert len(results) == 2
|
|
assert all(r["run_id"] == "r1" for r in results)
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_by_thread(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
await repo.create(run_id="r1", thread_id="t1", rating=1)
|
|
await repo.create(run_id="r2", thread_id="t1", rating=-1)
|
|
await repo.create(run_id="r3", thread_id="t2", rating=1)
|
|
results = await repo.list_by_thread("t1")
|
|
assert len(results) == 2
|
|
assert all(r["thread_id"] == "t1" for r in results)
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
created = await repo.create(run_id="r1", thread_id="t1", rating=1)
|
|
deleted = await repo.delete(created["feedback_id"])
|
|
assert deleted is True
|
|
assert await repo.get(created["feedback_id"]) is None
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_nonexistent(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
deleted = await repo.delete("nonexistent")
|
|
assert deleted is False
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_aggregate_by_run(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
await repo.create(run_id="r1", thread_id="t1", rating=1, user_id="user-1")
|
|
await repo.create(run_id="r1", thread_id="t1", rating=1, user_id="user-2")
|
|
await repo.create(run_id="r1", thread_id="t1", rating=-1, user_id="user-3")
|
|
stats = await repo.aggregate_by_run("t1", "r1")
|
|
assert stats["total"] == 3
|
|
assert stats["positive"] == 2
|
|
assert stats["negative"] == 1
|
|
assert stats["run_id"] == "r1"
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_aggregate_empty(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
stats = await repo.aggregate_by_run("t1", "r1")
|
|
assert stats["total"] == 0
|
|
assert stats["positive"] == 0
|
|
assert stats["negative"] == 0
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_creates_new(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
record = await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
assert record["rating"] == 1
|
|
assert record["feedback_id"]
|
|
assert record["user_id"] == "u1"
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_updates_existing(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
first = await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
second = await repo.upsert(run_id="r1", thread_id="t1", rating=-1, user_id="u1", comment="changed my mind")
|
|
assert second["feedback_id"] == first["feedback_id"]
|
|
assert second["rating"] == -1
|
|
assert second["comment"] == "changed my mind"
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_retries_as_update_when_concurrent_insert_wins(self, tmp_path):
|
|
"""A losing first-submission INSERT updates the winner's row instead of raising IntegrityError."""
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
try:
|
|
first = await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
real_factory = repo._sf
|
|
repo._sf, _ = miss_select_factory(real_factory)
|
|
try:
|
|
second = await repo.upsert(run_id="r1", thread_id="t1", rating=-1, user_id="u1", comment="second tab")
|
|
finally:
|
|
repo._sf = real_factory
|
|
|
|
assert second["feedback_id"] == first["feedback_id"]
|
|
assert second["rating"] == -1
|
|
assert second["comment"] == "second tab"
|
|
rows = await repo.list_by_run("t1", "r1", user_id="u1")
|
|
assert [row["feedback_id"] for row in rows] == [first["feedback_id"]]
|
|
finally:
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_raises_the_integrity_error_once_retries_are_exhausted(self, tmp_path):
|
|
"""If every attempt still collides, the original IntegrityError propagates instead of being swallowed."""
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
try:
|
|
first = await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
real_factory = repo._sf
|
|
repo._sf, sessions = miss_select_factory(real_factory, misses=None)
|
|
try:
|
|
with pytest.raises(IntegrityError, match="feedback.thread_id, feedback.run_id, feedback.user_id"):
|
|
await repo.upsert(run_id="r1", thread_id="t1", rating=-1, user_id="u1")
|
|
finally:
|
|
repo._sf = real_factory
|
|
|
|
assert [session.missed for session in sessions] == [_UPSERT_MAX_ATTEMPTS]
|
|
rows = await repo.list_by_run("t1", "r1", user_id="u1")
|
|
assert [(row["feedback_id"], row["rating"]) for row in rows] == [(first["feedback_id"], 1)]
|
|
finally:
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_different_users_separate(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
r1 = await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
r2 = await repo.upsert(run_id="r1", thread_id="t1", rating=-1, user_id="u2")
|
|
assert r1["feedback_id"] != r2["feedback_id"]
|
|
assert r1["rating"] == 1
|
|
assert r2["rating"] == -1
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_upsert_invalid_rating(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
with pytest.raises(ValueError):
|
|
await repo.upsert(run_id="r1", thread_id="t1", rating=0, user_id="u1")
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_by_run(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
deleted = await repo.delete_by_run(thread_id="t1", run_id="r1", user_id="u1")
|
|
assert deleted is True
|
|
results = await repo.list_by_run("t1", "r1", user_id="u1")
|
|
assert len(results) == 0
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_by_run_nonexistent(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
deleted = await repo.delete_by_run(thread_id="t1", run_id="r1", user_id="u1")
|
|
assert deleted is False
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_by_thread_grouped(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
await repo.upsert(run_id="r2", thread_id="t1", rating=-1, user_id="u1")
|
|
await repo.upsert(run_id="r3", thread_id="t2", rating=1, user_id="u1")
|
|
grouped = await repo.list_by_thread_grouped("t1", user_id="u1")
|
|
assert "r1" in grouped
|
|
assert "r2" in grouped
|
|
assert "r3" not in grouped
|
|
assert grouped["r1"]["rating"] == 1
|
|
assert grouped["r2"]["rating"] == -1
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_by_thread_grouped_empty(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
grouped = await repo.list_by_thread_grouped("t1", user_id="u1")
|
|
assert grouped == {}
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_by_run_ids_is_thread_and_owner_scoped(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
await repo.upsert(run_id="r2", thread_id="t1", rating=-1, user_id="u1")
|
|
await repo.upsert(run_id="r3", thread_id="t1", rating=1, user_id="u1")
|
|
await repo.upsert(run_id="r1", thread_id="t1", rating=-1, user_id="u2")
|
|
await repo.upsert(run_id="r2", thread_id="t2", rating=1, user_id="u1")
|
|
|
|
grouped = await repo.list_by_run_ids("t1", {"r1", "r2"}, user_id="u1")
|
|
|
|
assert set(grouped) == {"r1", "r2"}
|
|
assert grouped["r1"]["rating"] == 1
|
|
assert grouped["r2"]["rating"] == -1
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_unfiltered_reads_collapse_multi_user_feedback_deterministically(self, tmp_path):
|
|
"""Explicit-None reads (internal callers) collapse same-run feedback rows.
|
|
|
|
Several browser users can hold feedback on one run of a shared
|
|
NULL-owner thread; the per-run collapse must be deterministic —
|
|
the most recently created feedback wins, ``feedback_id`` breaks ties.
|
|
"""
|
|
import asyncio
|
|
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
await repo.upsert(run_id="r1", thread_id="t1", rating=1, user_id="u1")
|
|
# Guarantee a strictly later created_at than the first row.
|
|
await asyncio.sleep(0.01)
|
|
second = await repo.upsert(run_id="r1", thread_id="t1", rating=-1, user_id="u2")
|
|
|
|
for _ in range(3):
|
|
grouped = await repo.list_by_thread_grouped("t1", user_id=None)
|
|
by_run_ids = await repo.list_by_run_ids("t1", {"r1"}, user_id=None)
|
|
assert grouped["r1"]["feedback_id"] == second["feedback_id"]
|
|
assert grouped["r1"]["rating"] == -1
|
|
assert grouped["r1"]["user_id"] == "u2"
|
|
assert by_run_ids["r1"]["feedback_id"] == second["feedback_id"]
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_by_run_ids_empty_skips_query(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
|
|
assert await repo.list_by_run_ids("t1", set(), user_id="u1") == {}
|
|
await _cleanup()
|
|
|
|
|
|
# -- Follow-up association --
|
|
|
|
|
|
class TestFollowUpAssociation:
|
|
@pytest.mark.anyio
|
|
async def test_run_records_follow_up_via_memory_store(self):
|
|
"""MemoryRunStore stores follow_up_to_run_id in kwargs."""
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
await store.put("r1", thread_id="t1", status="success")
|
|
# MemoryRunStore doesn't have follow_up_to_run_id as a top-level param,
|
|
# but it can be passed via metadata
|
|
await store.put("r2", thread_id="t1", metadata={"follow_up_to_run_id": "r1"})
|
|
run = await store.get("r2")
|
|
assert run["metadata"]["follow_up_to_run_id"] == "r1"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_human_message_has_follow_up_metadata(self):
|
|
"""human_message event metadata includes follow_up_to_run_id."""
|
|
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
|
|
|
event_store = MemoryRunEventStore()
|
|
await event_store.put(
|
|
thread_id="t1",
|
|
run_id="r2",
|
|
event_type="human_message",
|
|
category="message",
|
|
content="Tell me more about that",
|
|
metadata={"follow_up_to_run_id": "r1"},
|
|
)
|
|
messages = await event_store.list_messages("t1")
|
|
assert messages[0]["metadata"]["follow_up_to_run_id"] == "r1"
|
|
|
|
@pytest.mark.anyio
|
|
async def test_follow_up_auto_detection_logic(self):
|
|
"""Simulate the auto-detection: latest successful run becomes follow_up_to."""
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
await store.put("r1", thread_id="t1", status="success")
|
|
await store.put("r2", thread_id="t1", status="error")
|
|
|
|
# Auto-detect: list_by_thread returns newest first
|
|
recent = await store.list_by_thread("t1", limit=1)
|
|
follow_up = None
|
|
if recent and recent[0].get("status") == "success":
|
|
follow_up = recent[0]["run_id"]
|
|
# r2 (error) is newest, so no follow_up detected
|
|
assert follow_up is None
|
|
|
|
# Now add a successful run
|
|
await store.put("r3", thread_id="t1", status="success")
|
|
recent = await store.list_by_thread("t1", limit=1)
|
|
follow_up = None
|
|
if recent and recent[0].get("status") == "success":
|
|
follow_up = recent[0]["run_id"]
|
|
assert follow_up == "r3"
|
|
|
|
|
|
class TestDeleteByThread:
|
|
"""Thread deletion clears feedback without crossing owner or thread scope."""
|
|
|
|
@pytest.mark.anyio
|
|
async def test_deletes_only_the_owners_rows_for_that_thread(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
try:
|
|
await repo.create(run_id="r1", thread_id="t1", rating=1, user_id="alice")
|
|
await repo.create(run_id="r2", thread_id="t1", rating=-1, user_id="alice")
|
|
await repo.create(run_id="r1", thread_id="t1", rating=1, user_id="bob")
|
|
await repo.create(run_id="r1", thread_id="t2", rating=1, user_id="alice")
|
|
|
|
count = await repo.delete_by_thread("t1", user_id="alice")
|
|
|
|
assert count == 2
|
|
assert await repo.list_by_thread("t1", user_id="alice") == []
|
|
assert len(await repo.list_by_thread("t1", user_id="bob")) == 1
|
|
assert len(await repo.list_by_thread("t2", user_id="alice")) == 1
|
|
finally:
|
|
await _cleanup()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_second_call_returns_zero(self, tmp_path):
|
|
repo = await _make_feedback_repo(tmp_path)
|
|
try:
|
|
await repo.create(run_id="r1", thread_id="t1", rating=1, user_id="alice")
|
|
|
|
assert await repo.delete_by_thread("t1", user_id="alice") == 1
|
|
assert await repo.delete_by_thread("t1", user_id="alice") == 0
|
|
finally:
|
|
await _cleanup()
|