mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-24 07:34:10 +08:00
fix(agent): scope chat history retrieval by flow_id to prevent cross-flow leak (#13087)
* fix(agent): scope chat history retrieval by flow_id to prevent cross-flow leak (#13059) AgentComponent.get_memory_data filtered chat history by session_id only, so the playground's default session names (e.g. "New Session 0") leaked conversation history across unrelated flows whenever two flows happened to use the same default name. The fix replaces the ad-hoc MemoryComponent spawn (whose internal _vertex is None and therefore cannot see the running flow's flow_id) with a direct aget_messages call that passes flow_id from the agent's own graph context. MemoryComponent.retrieve_messages is also updated to scope by self.graph.flow_id when available, so the standalone Message History component does not have the same leak. Adds regression tests covering: UUID coercion of string flow_ids, isolation of two flows that share a session name, graceful fallback when flow_id is missing or non-UUID, untouched external-memory path, and filtering out the current input message. * fix(agent): preserve n_messages=0 disable-memory contract and apply scoping to Cuga Review follow-up on #13087: - AgentComponent.get_memory_data regressed n_messages=0: before this PR the call went through MemoryComponent.retrieve_messages which short-circuits to [] when n_messages==0. The direct aget_messages call lost that contract, because `if self.n_messages:` is falsy for 0 and `messages[-0:]` returns every message fetched under limit=10000. - CugaComponent.get_memory_data had the same leak pattern (issue #13059) because it also spawned an ad-hoc MemoryComponent whose _vertex is None. Extract aget_agent_chat_history(session_id, flow_id, context_id, n_messages) into memory.py: * returns [] when n_messages == 0 * coerces flow_id to UUID (gracefully falls back on invalid values) * applies n_messages slicing Both AgentComponent and CugaComponent now route through it. Tests: * aget_agent_chat_history: passes flow_id as UUID, short-circuits on n_messages=0, slices to most-recent N, returns all on n_messages=None, falls back to unscoped query on invalid flow_id. * AgentComponent integration: routes through helper with flow_id, filters out current input, n_messages=0 disables memory (asserts aget_messages is never awaited so a future regression resurfaces as a real DB call). * CugaComponent integration: scopes by flow_id, n_messages=0 disables memory. Module-import gated for envs without optional Cuga deps. * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes (attempt 2/3) * [autofix.ci] apply automated fixes * Template update * [autofix.ci] apply automated fixes * Update .secrets.baseline * [autofix.ci] apply automated fixes * Update starter projects * [autofix.ci] apply automated fixes * Update .secrets.baseline * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes * fix(memory): address PR #13087 review — symmetric flow_id, DESC fetch, loud fallback Addresses the review comments on PR #13087: I1 (symmetric flow_id handling): `store_message` now routes the internal-memory write through `_coerce_flow_id_to_uuid(_safe_graph_flow_id(self))` instead of accessing `self.graph.flow_id` directly, so an ad-hoc / custom-subclass MemoryComponent without `_vertex` no longer crashes on the write half before reaching the safe read half. Both calls share a single `flow_id_scope` variable. I2 (>10k row ordering bug): both `aget_agent_chat_history` and `MemoryComponent.retrieve_messages` (internal-memory branch) now query in `DESC` order with `limit=n_messages` (falling back to `MAX_CHAT_HISTORY_FETCH_LIMIT` when no explicit limit is set), then reverse to the caller's preferred order. The previous ASC + slice shape returned the chronological FIRST window once sessions exceeded 10k rows, silently serving stale history. I3 (loud fallback observability): when `_coerce_flow_id_to_uuid` is forced to return `None` for a malformed `flow_id`, the log call is now `logger.error` (was `warning`) and carries a structured `event` tag (`memory_flow_id_unscoped`). The unbounded-fetch ceiling-hit path emits a parallel `memory_chat_history_limit_reached` warning. Both events are explicit alert hooks for observability pipelines so a regression cannot silently re-enable the cross-flow leak that motivated the PR. R1 (magic number): extracted `MAX_CHAT_HISTORY_FETCH_LIMIT = 10_000` as a module-level constant; reused in both call sites and the test suite. R2 (signature): tightened `_coerce_flow_id_to_uuid` and `_safe_graph_flow_id` from `Any` to `str | UUID | None`; tightened `aget_agent_chat_history`'s `flow_id` argument accordingly. Tests: extended `test_memory_flow_id_scoping.py` with coverage for each of the above — symmetric store/read flow_id (I1), DESC+reverse ordering on retrieve_messages (I2), ceiling-warning emission (I2), structured error log on flow_id fallback (I3), and explicit-limit suppression of the warning. All 26 tests pass. * [autofix.ci] apply automated fixes * Template update --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Eric Hare
@ -2438,7 +2438,7 @@
|
||||
"filename": "src/backend/base/langflow/initial_setup/starter_projects/Nvidia Remix.json",
|
||||
"hashed_secret": "b8205f6293ceac2bf593580250e865708c8a9aeb",
|
||||
"is_verified": false,
|
||||
"line_number": 2136
|
||||
"line_number": 2157
|
||||
}
|
||||
],
|
||||
"src/backend/base/langflow/initial_setup/starter_projects/Pok\u00e9dex Agent.json": [
|
||||
@ -2833,14 +2833,14 @@
|
||||
"filename": "src/backend/base/langflow/initial_setup/starter_projects/Youtube Analysis.json",
|
||||
"hashed_secret": "d6e6d7b4b115cd3b9d172623199f8c403055fecc",
|
||||
"is_verified": false,
|
||||
"line_number": 1323
|
||||
"line_number": 1344
|
||||
},
|
||||
{
|
||||
"type": "Hex High Entropy String",
|
||||
"filename": "src/backend/base/langflow/initial_setup/starter_projects/Youtube Analysis.json",
|
||||
"hashed_secret": "54ed260e3bc31bc77ee06754dff850981d39a66c",
|
||||
"is_verified": false,
|
||||
"line_number": 2095,
|
||||
"line_number": 2116,
|
||||
"is_secret": false
|
||||
}
|
||||
],
|
||||
@ -9031,5 +9031,5 @@
|
||||
}
|
||||
]
|
||||
},
|
||||
"generated_at": "2026-05-19T17:11:47Z"
|
||||
"generated_at": "2026-05-20T18:39:19Z"
|
||||
}
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@ -0,0 +1,540 @@
|
||||
"""Tests that MemoryComponent scopes chat history retrieval by flow_id.
|
||||
|
||||
Regression test for https://github.com/langflow-ai/langflow/issues/13059
|
||||
|
||||
Before the fix, ``MemoryComponent.retrieve_messages`` called
|
||||
``aget_messages`` without ``flow_id``. Because the Langflow playground assigns
|
||||
default session names (e.g. "New Session 0") that are not unique across flows,
|
||||
this caused chat history from Flow A to leak into Flow B whenever both used
|
||||
the same session name. The Agent component is affected because it delegates
|
||||
its chat-history fetch to this method.
|
||||
|
||||
These tests assert that the call into ``aget_messages`` is now scoped by the
|
||||
graph's ``flow_id`` (coerced to ``UUID``) so flows can no longer cross-read
|
||||
each other's history.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import pytest
|
||||
from lfx.components.models_and_agents.memory import (
|
||||
MAX_CHAT_HISTORY_FETCH_LIMIT,
|
||||
MemoryComponent,
|
||||
_coerce_flow_id_to_uuid,
|
||||
aget_agent_chat_history,
|
||||
)
|
||||
|
||||
|
||||
def _build_component(
|
||||
flow_id: str | UUID | None,
|
||||
session_id: str = "session-shared",
|
||||
) -> MemoryComponent:
|
||||
component = MemoryComponent()
|
||||
component.set(session_id=session_id, n_messages=10, order="Ascending")
|
||||
# ``Component.graph`` is a read-only property -> ``self._vertex.graph``,
|
||||
# so we wire the underlying ``_vertex`` instead of assigning ``graph``.
|
||||
component._vertex = SimpleNamespace(graph=SimpleNamespace(flow_id=flow_id, session_id=session_id))
|
||||
return component
|
||||
|
||||
|
||||
class TestCoerceFlowIdToUuid:
|
||||
"""Validate the small UUID-coercion helper used by retrieve_messages."""
|
||||
|
||||
def test_returns_none_for_missing_values(self):
|
||||
assert _coerce_flow_id_to_uuid(None) is None
|
||||
assert _coerce_flow_id_to_uuid("") is None
|
||||
|
||||
def test_passes_uuid_through_unchanged(self):
|
||||
value = uuid4()
|
||||
assert _coerce_flow_id_to_uuid(value) is value
|
||||
|
||||
def test_parses_uuid_string(self):
|
||||
raw = "11111111-1111-1111-1111-111111111111"
|
||||
assert _coerce_flow_id_to_uuid(raw) == UUID(raw)
|
||||
|
||||
def test_returns_none_for_invalid_string_and_does_not_raise(self):
|
||||
# Synthetic/test flow IDs may not be UUIDs; we must degrade gracefully.
|
||||
assert _coerce_flow_id_to_uuid("not-a-uuid") is None
|
||||
|
||||
def test_invalid_flow_id_emits_structured_error_log(self):
|
||||
"""Fallback to unscoped retrieval is a privacy regression — alert on it.
|
||||
|
||||
See PR #13087 review I3: this path re-enables the cross-flow leak
|
||||
that issue #13059 closed. The log call uses ``error`` level + a
|
||||
structured ``event`` tag so observability pipelines can alert.
|
||||
"""
|
||||
with patch("lfx.components.models_and_agents.memory.logger") as mock_logger:
|
||||
_coerce_flow_id_to_uuid("not-a-uuid")
|
||||
|
||||
mock_logger.error.assert_called_once()
|
||||
call = mock_logger.error.call_args
|
||||
assert "memory_flow_id_unscoped" in call.args[0]
|
||||
assert call.kwargs.get("extra", {}).get("event") == "memory_flow_id_unscoped"
|
||||
|
||||
|
||||
class TestRetrieveMessagesPassesFlowId:
|
||||
"""The leak fix: retrieve_messages must pass flow_id to aget_messages."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passes_graph_flow_id_as_uuid(self):
|
||||
flow_id_str = "22222222-2222-2222-2222-222222222222"
|
||||
component = _build_component(flow_id=flow_id_str)
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_get:
|
||||
await component.retrieve_messages()
|
||||
|
||||
mock_get.assert_awaited_once()
|
||||
kwargs = mock_get.await_args.kwargs
|
||||
assert kwargs["flow_id"] == UUID(flow_id_str), (
|
||||
"retrieve_messages must scope by flow_id so default session names "
|
||||
"do not leak chat history across flows (issue #13059)."
|
||||
)
|
||||
assert kwargs["session_id"] == "session-shared"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_isolates_two_flows_sharing_a_session_name(self):
|
||||
"""Reproduction of the leak: same session name, different flow_ids.
|
||||
|
||||
Flow A and Flow B both use session ``New Session 0``. With the fix in
|
||||
place, the call into ``aget_messages`` is scoped by each flow's UUID,
|
||||
so the database query can no longer return cross-flow rows.
|
||||
"""
|
||||
flow_a = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
|
||||
flow_b = "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"
|
||||
shared_session = "New Session 0"
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_get:
|
||||
await _build_component(flow_id=flow_a, session_id=shared_session).retrieve_messages()
|
||||
await _build_component(flow_id=flow_b, session_id=shared_session).retrieve_messages()
|
||||
|
||||
assert mock_get.await_count == 2
|
||||
first_call_kwargs = mock_get.await_args_list[0].kwargs
|
||||
second_call_kwargs = mock_get.await_args_list[1].kwargs
|
||||
assert first_call_kwargs["flow_id"] == UUID(flow_a)
|
||||
assert second_call_kwargs["flow_id"] == UUID(flow_b)
|
||||
# Same session name reaches the DB layer, but flow_id now disambiguates.
|
||||
assert first_call_kwargs["session_id"] == shared_session
|
||||
assert second_call_kwargs["session_id"] == shared_session
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_falls_back_to_none_when_flow_id_missing(self):
|
||||
component = _build_component(flow_id=None)
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_get:
|
||||
await component.retrieve_messages()
|
||||
|
||||
assert mock_get.await_args.kwargs["flow_id"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_falls_back_to_none_when_flow_id_is_not_a_uuid(self):
|
||||
"""Tests with synthetic graph IDs must not crash retrieval."""
|
||||
component = _build_component(flow_id="not-a-uuid")
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_get:
|
||||
await component.retrieve_messages()
|
||||
|
||||
assert mock_get.await_args.kwargs["flow_id"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_external_memory_path_is_untouched(self):
|
||||
"""External Memory providers have no concept of flow_id; do not pass it."""
|
||||
|
||||
class _FakeExternalMemory:
|
||||
session_id: str | None = None
|
||||
context_id: str | None = None
|
||||
|
||||
async def aget_messages(self):
|
||||
return []
|
||||
|
||||
component = _build_component(flow_id="33333333-3333-3333-3333-333333333333")
|
||||
component.set(memory=_FakeExternalMemory())
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_get:
|
||||
await component.retrieve_messages()
|
||||
|
||||
mock_get.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetches_desc_and_reverses_for_ascending_order(self):
|
||||
"""PR #13087 review I2: avoid the >10k row ASC slice trap.
|
||||
|
||||
Querying ``order=ASC, limit=MAX`` then slicing ``[-n:]`` would return
|
||||
the chronological FIRST window on huge sessions. We query DESC with
|
||||
the actual limit and reverse to the user-selected order, so the slice
|
||||
is always the genuine most-recent messages.
|
||||
"""
|
||||
component = _build_component(flow_id="22222222-2222-2222-2222-222222222222")
|
||||
# MemoryComponent ``order`` defaults to Ascending in _build_component.
|
||||
# The mock returns DESC-ordered data (most recent first), as the DB would.
|
||||
desc_rows = [SimpleNamespace(id=f"msg-{i}") for i in (4, 3, 2, 1, 0)]
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=desc_rows),
|
||||
) as mock_get:
|
||||
result = await component.retrieve_messages()
|
||||
|
||||
assert mock_get.await_args.kwargs["order"] == "DESC"
|
||||
assert mock_get.await_args.kwargs["limit"] == 10 # _build_component sets n_messages=10
|
||||
assert [m.id for m in result] == ["msg-0", "msg-1", "msg-2", "msg-3", "msg-4"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_descending_order_passthrough_keeps_desc_view(self):
|
||||
component = _build_component(flow_id="22222222-2222-2222-2222-222222222222")
|
||||
component.set(order="Descending")
|
||||
desc_rows = [SimpleNamespace(id=f"msg-{i}") for i in (4, 3, 2, 1, 0)]
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=desc_rows),
|
||||
):
|
||||
result = await component.retrieve_messages()
|
||||
|
||||
assert [m.id for m in result] == ["msg-4", "msg-3", "msg-2", "msg-1", "msg-0"]
|
||||
|
||||
|
||||
class TestStoreMessagePassesFlowId:
|
||||
"""PR #13087 review I1: write/read flow_id handling must be symmetric.
|
||||
|
||||
The original PR protected ``retrieve_messages`` against ``self.graph``
|
||||
raising on ad-hoc instantiation but left ``store_message`` calling
|
||||
``self.graph.flow_id`` directly — so a custom subclass or integration
|
||||
test would crash on the write path before reaching the safe read path.
|
||||
Both paths now go through ``_coerce_flow_id_to_uuid(_safe_graph_flow_id(self))``.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_passes_coerced_flow_id_to_astore_and_aget(self):
|
||||
flow_id_str = "22222222-2222-2222-2222-222222222222"
|
||||
component = _build_component(flow_id=flow_id_str)
|
||||
component.set(mode="Store", message="hi", sender="User", sender_name="Tester")
|
||||
|
||||
stored = SimpleNamespace(id="stored-1")
|
||||
with (
|
||||
patch(
|
||||
"lfx.components.models_and_agents.memory.astore_message",
|
||||
new=AsyncMock(return_value=None),
|
||||
) as mock_store,
|
||||
patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[stored]),
|
||||
) as mock_get,
|
||||
):
|
||||
await component.store_message()
|
||||
|
||||
assert mock_store.await_args.kwargs["flow_id"] == UUID(flow_id_str)
|
||||
assert mock_get.await_args.kwargs["flow_id"] == UUID(flow_id_str)
|
||||
# Symmetric scope: the same coerced UUID is used for both calls.
|
||||
assert mock_store.await_args.kwargs["flow_id"] == mock_get.await_args.kwargs["flow_id"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_does_not_crash_without_vertex(self):
|
||||
"""When _vertex is None, ``self.graph`` raises AttributeError.
|
||||
|
||||
Pre-fix, ``store_message`` accessed ``self.graph.flow_id`` directly
|
||||
and crashed before reaching the safe read path. Now both paths
|
||||
share the same defensive lookup.
|
||||
"""
|
||||
component = MemoryComponent()
|
||||
component.set(
|
||||
session_id="s",
|
||||
n_messages=10,
|
||||
order="Ascending",
|
||||
mode="Store",
|
||||
message="hi",
|
||||
sender="User",
|
||||
sender_name="Tester",
|
||||
)
|
||||
# No _vertex assigned — ``self.graph`` will raise AttributeError.
|
||||
stored = SimpleNamespace(id="stored-1")
|
||||
with (
|
||||
patch(
|
||||
"lfx.components.models_and_agents.memory.astore_message",
|
||||
new=AsyncMock(return_value=None),
|
||||
) as mock_store,
|
||||
patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[stored]),
|
||||
),
|
||||
):
|
||||
await component.store_message()
|
||||
|
||||
assert mock_store.await_args.kwargs["flow_id"] is None
|
||||
|
||||
|
||||
class TestAgetAgentChatHistoryHelper:
|
||||
"""The shared helper centralizes the agent-side memory contract.
|
||||
|
||||
Both ``AgentComponent`` and ``CugaComponent`` route through this helper,
|
||||
so testing the helper directly covers their common behavior.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passes_flow_id_as_uuid(self):
|
||||
flow_id_str = "44444444-4444-4444-4444-444444444444"
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_get:
|
||||
await aget_agent_chat_history(
|
||||
session_id="New Session 0",
|
||||
flow_id=flow_id_str,
|
||||
context_id="",
|
||||
n_messages=10,
|
||||
)
|
||||
|
||||
mock_get.assert_awaited_once()
|
||||
kwargs = mock_get.await_args.kwargs
|
||||
assert kwargs["flow_id"] == UUID(flow_id_str), "aget_agent_chat_history must scope by flow_id (issue #13059)."
|
||||
assert kwargs["session_id"] == "New Session 0"
|
||||
# PR #13087 review I2: query DESC + limit=n_messages so sessions
|
||||
# exceeding MAX_CHAT_HISTORY_FETCH_LIMIT still get the real tail.
|
||||
assert kwargs["order"] == "DESC"
|
||||
assert kwargs["limit"] == 10
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_n_messages_zero_short_circuits_without_querying(self):
|
||||
"""Regression: ``n_messages == 0`` means "memory disabled".
|
||||
|
||||
Before this short-circuit, ``messages[-0:]`` returned the full
|
||||
bounded result, so users who set the field to 0 to disable chat
|
||||
memory unexpectedly got full history back.
|
||||
"""
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[SimpleNamespace(id=f"msg-{i}") for i in range(5)]),
|
||||
) as mock_get:
|
||||
result = await aget_agent_chat_history(
|
||||
session_id="s",
|
||||
flow_id="66666666-6666-6666-6666-666666666666",
|
||||
n_messages=0,
|
||||
)
|
||||
|
||||
assert result == []
|
||||
mock_get.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slices_to_n_messages_most_recent(self):
|
||||
# DB returns DESC (most recent first); helper reverses to ASC.
|
||||
desc_rows = [SimpleNamespace(id=f"msg-{i}") for i in (4, 3)]
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=desc_rows),
|
||||
) as mock_get:
|
||||
result = await aget_agent_chat_history(
|
||||
session_id="s",
|
||||
flow_id="77777777-7777-7777-7777-777777777777",
|
||||
n_messages=2,
|
||||
)
|
||||
|
||||
assert mock_get.await_args.kwargs["limit"] == 2
|
||||
assert mock_get.await_args.kwargs["order"] == "DESC"
|
||||
assert [m.id for m in result] == ["msg-3", "msg-4"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_n_messages_returns_all_fetched(self):
|
||||
# DB returns DESC; helper reverses for the agent prompt.
|
||||
desc_rows = [SimpleNamespace(id=f"msg-{i}") for i in (2, 1, 0)]
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=desc_rows),
|
||||
) as mock_get:
|
||||
result = await aget_agent_chat_history(
|
||||
session_id="s",
|
||||
flow_id=None,
|
||||
n_messages=None,
|
||||
)
|
||||
|
||||
# No explicit limit ⇒ fall back to MAX_CHAT_HISTORY_FETCH_LIMIT.
|
||||
assert mock_get.await_args.kwargs["limit"] == MAX_CHAT_HISTORY_FETCH_LIMIT
|
||||
assert [m.id for m in result] == ["msg-0", "msg-1", "msg-2"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_flow_id_falls_back_to_unscoped_query(self):
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_get:
|
||||
await aget_agent_chat_history(
|
||||
session_id="s",
|
||||
flow_id="not-a-uuid",
|
||||
n_messages=10,
|
||||
)
|
||||
|
||||
assert mock_get.await_args.kwargs["flow_id"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logs_warning_when_unbounded_fetch_hits_ceiling(self):
|
||||
"""PR #13087 review I2: surface the silent-truncation case.
|
||||
|
||||
When the caller did not pass ``n_messages`` and the DB returned
|
||||
exactly ``MAX_CHAT_HISTORY_FETCH_LIMIT`` rows, older history may
|
||||
have been silently dropped. Emit a structured warning so on-call
|
||||
has something to grep when a user reports forgotten context.
|
||||
"""
|
||||
# Simulate the DB returning exactly the ceiling — older rows truncated.
|
||||
rows = [SimpleNamespace(id=f"msg-{i}") for i in range(MAX_CHAT_HISTORY_FETCH_LIMIT)]
|
||||
with (
|
||||
patch("lfx.components.models_and_agents.memory.logger") as mock_logger,
|
||||
patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=rows),
|
||||
),
|
||||
):
|
||||
await aget_agent_chat_history(
|
||||
session_id="s",
|
||||
flow_id=None,
|
||||
n_messages=None,
|
||||
)
|
||||
|
||||
mock_logger.warning.assert_called_once()
|
||||
call = mock_logger.warning.call_args
|
||||
assert "memory_chat_history_limit_reached" in call.args[0]
|
||||
assert call.kwargs.get("extra", {}).get("event") == "memory_chat_history_limit_reached"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_n_messages_at_ceiling_does_not_warn(self):
|
||||
"""The ceiling warning is for ``n_messages=None`` only — explicit limits are honored."""
|
||||
rows = [SimpleNamespace(id=f"msg-{i}") for i in range(MAX_CHAT_HISTORY_FETCH_LIMIT)]
|
||||
with (
|
||||
patch("lfx.components.models_and_agents.memory.logger") as mock_logger,
|
||||
patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=rows),
|
||||
),
|
||||
):
|
||||
await aget_agent_chat_history(
|
||||
session_id="s",
|
||||
flow_id=None,
|
||||
n_messages=MAX_CHAT_HISTORY_FETCH_LIMIT,
|
||||
)
|
||||
|
||||
mock_logger.warning.assert_not_called()
|
||||
|
||||
|
||||
class TestAgentGetMemoryDataIntegration:
|
||||
"""End-to-end checks at the AgentComponent boundary."""
|
||||
|
||||
@staticmethod
|
||||
def _make_agent(flow_id: str | UUID | None, session_id: str = "New Session 0", n_messages: int = 10):
|
||||
from lfx.components.models_and_agents.agent import AgentComponent
|
||||
|
||||
agent = AgentComponent.__new__(AgentComponent)
|
||||
agent._vertex = SimpleNamespace(graph=SimpleNamespace(flow_id=flow_id, session_id=session_id))
|
||||
agent.context_id = ""
|
||||
agent.n_messages = n_messages
|
||||
agent.input_value = SimpleNamespace(id="current-input-id")
|
||||
return agent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_routes_through_helper_with_flow_id(self):
|
||||
flow_id_str = "88888888-8888-8888-8888-888888888888"
|
||||
agent = self._make_agent(flow_id=flow_id_str)
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.agent.aget_agent_chat_history",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_helper:
|
||||
await agent.get_memory_data()
|
||||
|
||||
mock_helper.assert_awaited_once()
|
||||
kwargs = mock_helper.await_args.kwargs
|
||||
assert kwargs["flow_id"] == flow_id_str
|
||||
assert kwargs["session_id"] == "New Session 0"
|
||||
assert kwargs["n_messages"] == 10
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_filters_out_current_input_message(self):
|
||||
"""The agent must not echo the current input back as chat history."""
|
||||
agent = self._make_agent(flow_id="55555555-5555-5555-5555-555555555555")
|
||||
|
||||
current = SimpleNamespace(id="current-input-id", text="ping")
|
||||
old = SimpleNamespace(id="old-msg-id", text="earlier")
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.agent.aget_agent_chat_history",
|
||||
new=AsyncMock(return_value=[old, current]),
|
||||
):
|
||||
result = await agent.get_memory_data()
|
||||
|
||||
assert [m.id for m in result] == ["old-msg-id"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_n_messages_zero_disables_memory(self):
|
||||
"""Regression: setting "Number of Chat History Messages" to 0 must disable memory."""
|
||||
agent = self._make_agent(flow_id="99999999-9999-9999-9999-999999999999", n_messages=0)
|
||||
|
||||
# Patch the underlying DB query so a regression resurfaces as a real call.
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[SimpleNamespace(id=f"msg-{i}") for i in range(5)]),
|
||||
) as mock_get:
|
||||
result = await agent.get_memory_data()
|
||||
|
||||
assert result == []
|
||||
mock_get.assert_not_awaited()
|
||||
|
||||
|
||||
class TestCugaGetMemoryDataIntegration:
|
||||
"""The Cuga agent had the same leak pattern; verify the fix reaches it too."""
|
||||
|
||||
@staticmethod
|
||||
def _make_cuga(flow_id: str | UUID | None, session_id: str = "shared", n_messages: int = 10):
|
||||
try:
|
||||
from lfx.components.cuga import cuga_agent
|
||||
except Exception as exc: # pragma: no cover - optional deps
|
||||
pytest.skip(f"cuga_agent not importable in this env: {exc}")
|
||||
|
||||
agent = cuga_agent.CugaComponent.__new__(cuga_agent.CugaComponent)
|
||||
agent._vertex = SimpleNamespace(graph=SimpleNamespace(flow_id=flow_id, session_id=session_id))
|
||||
agent.n_messages = n_messages
|
||||
agent.input_value = SimpleNamespace(id="current-input-id")
|
||||
return agent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cuga_scopes_by_flow_id(self):
|
||||
flow_id_str = "cccccccc-cccc-cccc-cccc-cccccccccccc"
|
||||
agent = self._make_cuga(flow_id=flow_id_str)
|
||||
|
||||
with patch(
|
||||
"lfx.components.cuga.cuga_agent.aget_agent_chat_history",
|
||||
new=AsyncMock(return_value=[]),
|
||||
) as mock_helper:
|
||||
await agent.get_memory_data()
|
||||
|
||||
kwargs = mock_helper.await_args.kwargs
|
||||
assert kwargs["flow_id"] == flow_id_str
|
||||
assert kwargs["session_id"] == "shared"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cuga_n_messages_zero_disables_memory(self):
|
||||
agent = self._make_cuga(flow_id="dddddddd-dddd-dddd-dddd-dddddddddddd", n_messages=0)
|
||||
|
||||
with patch(
|
||||
"lfx.components.models_and_agents.memory.aget_messages",
|
||||
new=AsyncMock(return_value=[SimpleNamespace(id=f"msg-{i}") for i in range(3)]),
|
||||
) as mock_get:
|
||||
result = await agent.get_memory_data()
|
||||
|
||||
assert result == []
|
||||
mock_get.assert_not_awaited()
|
||||
File diff suppressed because one or more lines are too long
@ -20,7 +20,7 @@ from lfx.base.models.model_input_constants import (
|
||||
from lfx.base.models.model_utils import get_model_name
|
||||
from lfx.components.helpers import CurrentDateComponent
|
||||
from lfx.components.langchain_utilities.tool_calling import ToolCallingAgentComponent
|
||||
from lfx.components.models_and_agents.memory import MemoryComponent
|
||||
from lfx.components.models_and_agents.memory import MemoryComponent, aget_agent_chat_history
|
||||
from lfx.custom.custom_component.component import _get_component_toolkit
|
||||
from lfx.custom.utils import update_component_build_config
|
||||
from lfx.field_typing import Tool
|
||||
@ -488,10 +488,13 @@ class CugaComponent(ToolCallingAgentComponent):
|
||||
logger.debug(f"[CUGA] input_value type: {type(self.input_value)}")
|
||||
logger.debug(f"[CUGA] input_value id: {getattr(self.input_value, 'id', None)}")
|
||||
|
||||
messages = (
|
||||
await MemoryComponent(**self.get_base_args())
|
||||
.set(session_id=str(self.graph.session_id), order="Ascending", n_messages=self.n_messages)
|
||||
.retrieve_messages()
|
||||
# Scope by flow_id (issue #13059): the ad-hoc MemoryComponent used previously
|
||||
# had no _vertex, so it could not see the running flow's flow_id and emitted
|
||||
# an unscoped query. The helper also honors n_messages == 0 as "disabled".
|
||||
messages = await aget_agent_chat_history(
|
||||
session_id=str(self.graph.session_id),
|
||||
flow_id=getattr(self.graph, "flow_id", None),
|
||||
n_messages=self.n_messages,
|
||||
)
|
||||
logger.debug(f"[CUGA] Retrieved {len(messages)} messages from memory")
|
||||
return [
|
||||
|
||||
@ -4,7 +4,7 @@ from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lfx.components.models_and_agents.memory import MemoryComponent
|
||||
from lfx.components.models_and_agents.memory import MemoryComponent, aget_agent_chat_history
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_core.tools import Tool
|
||||
@ -491,16 +491,15 @@ class AgentComponent(ToolCallingAgentComponent):
|
||||
return Data(data={"content": "", "error": str(exc)})
|
||||
|
||||
async def get_memory_data(self):
|
||||
# TODO: This is a temporary fix to avoid message duplication. We should develop a function for this.
|
||||
messages = (
|
||||
await MemoryComponent(**self.get_base_args())
|
||||
.set(
|
||||
session_id=self.graph.session_id,
|
||||
context_id=self.context_id,
|
||||
order="Ascending",
|
||||
n_messages=self.n_messages,
|
||||
)
|
||||
.retrieve_messages()
|
||||
# Scope by flow_id so default playground session names (e.g. "New Session 0")
|
||||
# cannot leak chat history across unrelated flows. See issue #13059.
|
||||
# The helper also returns [] when n_messages == 0, preserving the
|
||||
# explicit "memory disabled" contract from MemoryComponent.retrieve_messages.
|
||||
messages = await aget_agent_chat_history(
|
||||
session_id=self.graph.session_id,
|
||||
flow_id=getattr(self.graph, "flow_id", None),
|
||||
context_id=self.context_id,
|
||||
n_messages=self.n_messages,
|
||||
)
|
||||
return [
|
||||
message for message in messages if getattr(message, "id", None) != getattr(self.input_value, "id", None)
|
||||
|
||||
@ -1,8 +1,10 @@
|
||||
from typing import Any, cast
|
||||
from uuid import UUID
|
||||
|
||||
from lfx.custom.custom_component.component import Component
|
||||
from lfx.helpers.data import data_to_text
|
||||
from lfx.inputs.inputs import DropdownInput, HandleInput, IntInput, MessageTextInput, MultilineInput, TabInput
|
||||
from lfx.log.logger import logger
|
||||
from lfx.memory import aget_messages, astore_message
|
||||
from lfx.schema.data import Data
|
||||
from lfx.schema.dataframe import DataFrame
|
||||
@ -12,6 +14,106 @@ from lfx.template.field.base import Output
|
||||
from lfx.utils.component_utils import set_current_fields, set_field_display
|
||||
from lfx.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_NAME_AI, MESSAGE_SENDER_USER
|
||||
|
||||
# Cap on rows we will pull from the DB when the caller has not supplied an
|
||||
# explicit ``n_messages`` limit. This is a guardrail against unbounded fetches
|
||||
# (DB pressure, memory) on sessions that accumulate huge histories. The
|
||||
# trade-off is that sessions with more rows than this cap will not surface
|
||||
# their oldest messages when no ``n_messages`` is set. See ``aget_messages``
|
||||
# usage sites below for how this interacts with ordering.
|
||||
MAX_CHAT_HISTORY_FETCH_LIMIT = 10_000
|
||||
|
||||
|
||||
def _coerce_flow_id_to_uuid(flow_id: str | UUID | None) -> UUID | None:
|
||||
"""Coerce a graph flow_id (typically str) to UUID for DB filtering.
|
||||
|
||||
Returns ``None`` when ``flow_id`` is missing or cannot be parsed. The
|
||||
caller then falls back to the previous **unscoped** retrieval, which
|
||||
re-introduces the cross-flow leak that motivated PR #13087. We emit a
|
||||
structured ``error`` log on that path (rather than ``warning``) so
|
||||
observability can alert on it — see issue #13059 / PR #13087.
|
||||
"""
|
||||
if flow_id is None or flow_id == "":
|
||||
return None
|
||||
if isinstance(flow_id, UUID):
|
||||
return flow_id
|
||||
try:
|
||||
return UUID(str(flow_id))
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
# Loud, structured signal — this path means chat history is being
|
||||
# served without flow-scoping, which is the privacy bug PR #13087
|
||||
# was created to close. Anything matching this event is a candidate
|
||||
# for an observability alert.
|
||||
logger.error(
|
||||
"memory_flow_id_unscoped: flow_id %r is not a valid UUID; "
|
||||
"chat history will NOT be scoped by flow_id. This re-enables "
|
||||
"cross-flow leakage (issue #13059) for this request.",
|
||||
flow_id,
|
||||
extra={"event": "memory_flow_id_unscoped", "flow_id_repr": repr(flow_id)},
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _safe_graph_flow_id(component: Component) -> str | UUID | None:
|
||||
"""Best-effort lookup of the component's graph flow_id.
|
||||
|
||||
``Component.graph`` is a property that reaches into ``self._vertex.graph``;
|
||||
when a MemoryComponent is constructed ad-hoc (e.g. by the Agent component
|
||||
via ``MemoryComponent(**self.get_base_args())``), ``_vertex`` is ``None`` and
|
||||
accessing the property raises ``AttributeError``. Swallow that here so
|
||||
retrieval falls back to the previous unscoped behavior rather than crashing.
|
||||
"""
|
||||
try:
|
||||
graph = component.graph
|
||||
except AttributeError:
|
||||
return None
|
||||
return getattr(graph, "flow_id", None)
|
||||
|
||||
|
||||
async def aget_agent_chat_history(
|
||||
*,
|
||||
session_id: str | UUID | None,
|
||||
flow_id: str | UUID | None,
|
||||
context_id: str | None = None,
|
||||
n_messages: int | None = None,
|
||||
) -> list[Message]:
|
||||
"""Fetch chat history for an agent, scoped to a single flow.
|
||||
|
||||
Centralizes the contract previously implemented by
|
||||
``MemoryComponent.retrieve_messages`` for agent callers:
|
||||
|
||||
* Returns ``[]`` when ``n_messages == 0`` (memory explicitly disabled).
|
||||
Without this short-circuit, the bounded query would still execute and
|
||||
the caller's ``messages[-0:]`` would return everything.
|
||||
* Scopes by ``flow_id`` (coerced to ``UUID``) so default playground
|
||||
session names cannot leak history across flows (issue #13059).
|
||||
* Returns up to ``n_messages`` most recent messages in ascending order.
|
||||
Queries in ``DESC`` order with ``limit=n_messages`` so sessions with
|
||||
more than ``MAX_CHAT_HISTORY_FETCH_LIMIT`` rows still see the genuine
|
||||
most-recent slice, not the chronological first window.
|
||||
"""
|
||||
if n_messages == 0:
|
||||
return []
|
||||
fetch_limit = n_messages if n_messages else MAX_CHAT_HISTORY_FETCH_LIMIT
|
||||
messages = await aget_messages(
|
||||
session_id=session_id,
|
||||
context_id=context_id,
|
||||
flow_id=_coerce_flow_id_to_uuid(flow_id),
|
||||
limit=fetch_limit,
|
||||
order="DESC",
|
||||
)
|
||||
if not n_messages and len(messages) == MAX_CHAT_HISTORY_FETCH_LIMIT:
|
||||
# We hit the unbounded-fetch ceiling. The caller will likely see
|
||||
# stale "most-recent" history. Flag this so on-call has something
|
||||
# to grep when a user reports forgotten context.
|
||||
logger.warning(
|
||||
"memory_chat_history_limit_reached: hit MAX_CHAT_HISTORY_FETCH_LIMIT=%d "
|
||||
"without an explicit n_messages; older messages were not returned.",
|
||||
MAX_CHAT_HISTORY_FETCH_LIMIT,
|
||||
extra={"event": "memory_chat_history_limit_reached"},
|
||||
)
|
||||
# ``aget_messages`` returned DESC; reverse to ASC for the agent's prompt.
|
||||
return list(reversed(messages))
|
||||
|
||||
|
||||
class MemoryComponent(Component):
|
||||
display_name = "Message History"
|
||||
@ -194,13 +296,18 @@ class MemoryComponent(Component):
|
||||
if message.sender:
|
||||
stored_messages = [m for m in stored_messages if m.sender == message.sender]
|
||||
else:
|
||||
await astore_message(message, flow_id=self.graph.flow_id)
|
||||
# Single coerced scope used for both the write and the read-back,
|
||||
# so a missing/ad-hoc ``_vertex`` cannot crash the write half while
|
||||
# the read half degrades gracefully. See PR #13087 review I1.
|
||||
flow_id_scope = _coerce_flow_id_to_uuid(_safe_graph_flow_id(self))
|
||||
await astore_message(message, flow_id=flow_id_scope)
|
||||
stored_messages = (
|
||||
await aget_messages(
|
||||
session_id=message.session_id,
|
||||
context_id=message.context_id,
|
||||
sender_name=message.sender_name,
|
||||
sender=message.sender,
|
||||
flow_id=flow_id_scope,
|
||||
)
|
||||
or []
|
||||
)
|
||||
@ -250,17 +357,34 @@ class MemoryComponent(Component):
|
||||
expected_type = MESSAGE_SENDER_AI if sender_type == MESSAGE_SENDER_AI else MESSAGE_SENDER_USER
|
||||
stored = [m for m in stored if m.type == expected_type]
|
||||
else:
|
||||
# For internal memory, we always fetch the last N messages by ordering by DESC
|
||||
# For internal memory, fetch the last N messages by ordering DESC at the
|
||||
# DB layer so sessions with more than ``MAX_CHAT_HISTORY_FETCH_LIMIT``
|
||||
# rows still return the genuine most-recent slice rather than the
|
||||
# chronological first window. See PR #13087 review I2.
|
||||
#
|
||||
# Scope by flow_id so default session names (e.g. "New Session 0") do not
|
||||
# leak chat history across unrelated flows. See issue #13059.
|
||||
flow_id_scope = _coerce_flow_id_to_uuid(_safe_graph_flow_id(self))
|
||||
fetch_limit = n_messages if n_messages else MAX_CHAT_HISTORY_FETCH_LIMIT
|
||||
stored = await aget_messages(
|
||||
sender=sender_type,
|
||||
sender_name=sender_name,
|
||||
session_id=session_id,
|
||||
context_id=context_id,
|
||||
limit=10000,
|
||||
order=order,
|
||||
flow_id=flow_id_scope,
|
||||
limit=fetch_limit,
|
||||
order="DESC",
|
||||
)
|
||||
if n_messages:
|
||||
stored = stored[-n_messages:] # Get last N messages
|
||||
if not n_messages and len(stored) == MAX_CHAT_HISTORY_FETCH_LIMIT:
|
||||
logger.warning(
|
||||
"memory_chat_history_limit_reached: hit MAX_CHAT_HISTORY_FETCH_LIMIT=%d "
|
||||
"without an explicit n_messages; older messages were not returned.",
|
||||
MAX_CHAT_HISTORY_FETCH_LIMIT,
|
||||
extra={"event": "memory_chat_history_limit_reached"},
|
||||
)
|
||||
# Honor the user-selected order: we fetched DESC, reverse if ASC requested.
|
||||
if order == "ASC":
|
||||
stored = list(reversed(stored))
|
||||
|
||||
# self.status = stored
|
||||
return cast("Data", stored)
|
||||
|
||||
Reference in New Issue
Block a user