diff --git a/backend/chainlit/data/dynamodb.py b/backend/chainlit/data/dynamodb.py index b79d2018cf..9b29c372e9 100644 --- a/backend/chainlit/data/dynamodb.py +++ b/backend/chainlit/data/dynamodb.py @@ -177,8 +177,8 @@ async def delete_feedback(self, feedback_id: str) -> bool: # feedback id = THREAD#{thread_id}::STEP#{step_id} thread_id, step_id = feedback_id.split("::") - thread_id = thread_id.strip("THREAD#") - step_id = step_id.strip("STEP#") + thread_id = thread_id.removeprefix("THREAD#") + step_id = step_id.removeprefix("STEP#") self.client.update_item( TableName=self.table_name, @@ -500,8 +500,8 @@ async def list_threads( for item in response["Items"]: deserialized_item: Dict[str, Any] = self._deserialize_item(item) thread = ThreadDict( # type: ignore - id=deserialized_item["PK"].strip("THREAD#"), - createdAt=deserialized_item["UserThreadSK"].strip("TS#"), + id=deserialized_item["PK"].removeprefix("THREAD#"), + createdAt=deserialized_item["UserThreadSK"].removeprefix("TS#"), name=deserialized_item["name"], ) paginated_response.data.append(thread) diff --git a/backend/tests/data/test_dynamodb.py b/backend/tests/data/test_dynamodb.py new file mode 100644 index 0000000000..f3462b0ffb --- /dev/null +++ b/backend/tests/data/test_dynamodb.py @@ -0,0 +1,57 @@ +"""Tests for the DynamoDB data layer's thread and step identifier handling. + +DynamoDB stores a thread as ``PK="THREAD#"`` and a step as +``SK="STEP#"``. Reading those identifiers back must strip the prefix +only, never the characters the identifier happens to end with. +""" + +from unittest.mock import Mock + +import pytest + +from chainlit.data.dynamodb import DynamoDBDataLayer +from chainlit.types import Pagination, ThreadFilter + + +def _make_layer(items=None): + client = Mock() + client.query.return_value = {"Items": items or []} + return DynamoDBDataLayer(table_name="test-table", client=client), client + + +@pytest.mark.asyncio +async def test_list_threads_keeps_thread_ids_ending_in_prefix_characters(): + """A thread id ending in T/H/R/E/A/D keeps every character. + + ``str.strip("THREAD#")`` treats its argument as a set of characters, so it + eats the id's own trailing letters; ``removeprefix`` removes the prefix only. + """ + layer, _ = _make_layer( + [ + { + "PK": {"S": "THREAD#DEMO"}, + "UserThreadSK": {"S": "TS#2026-01-01T00:00:00"}, + "name": {"S": "demo thread"}, + } + ] + ) + + response = await layer.list_threads( + Pagination(first=10), ThreadFilter(userId="user-1") + ) + + assert [thread["id"] for thread in response.data] == ["DEMO"] + assert response.data[0]["createdAt"] == "2026-01-01T00:00:00" + + +@pytest.mark.asyncio +async def test_delete_feedback_builds_the_key_from_the_full_identifiers(): + """The delete key must reuse the ids from the feedback id verbatim.""" + layer, client = _make_layer() + + await layer.delete_feedback("THREAD#CODE::STEP#SEED") + + assert client.update_item.call_args.kwargs["Key"] == { + "PK": {"S": "THREAD#CODE"}, + "SK": {"S": "STEP#SEED"}, + }