|
| 1 | +"""Regression tests: ambient flow-scope defaults ``aget_messages`` by ``flow_id`` (issue #13059). |
| 2 | +
|
| 3 | +Langflow executes the *frozen* component ``code`` embedded in each saved flow, not the installed |
| 4 | +library version (``lfx.interface.initialize.loading.instantiate_class`` -> ``eval_custom_component_code``). |
| 5 | +A flow saved before PR #13087 therefore carries the old ``retrieve_messages`` that calls |
| 6 | +``aget_messages`` WITHOUT ``flow_id`` and leaks chat history across flows on a colliding |
| 7 | +``session_id`` — even when running on a patched server. |
| 8 | +
|
| 9 | +Defense-in-depth: the engine binds the executing graph's ``flow_id`` to a ContextVar so |
| 10 | +``aget_messages`` can default the scope. This closes the leak for old frozen flows without: |
| 11 | +- touching frozen code (impossible), |
| 12 | +- changing callers that pass ``flow_id`` explicitly (the default only applies when it is ``None``), |
| 13 | +- changing callers outside a graph run (the ContextVar is unset -> identical to today). |
| 14 | +""" |
| 15 | + |
| 16 | +from typing import cast |
| 17 | +from uuid import uuid4 |
| 18 | + |
| 19 | +from langflow.memory import aadd_messages, aget_messages |
| 20 | +from langflow.schema.message import Message |
| 21 | +from lfx.components.input_output import ChatOutput |
| 22 | +from lfx.components.models_and_agents.memory import MemoryComponent |
| 23 | +from lfx.graph.graph.base import Graph |
| 24 | +from lfx.memory.flow_context import reset_current_flow_id, set_current_flow_id |
| 25 | +from lfx.schema.data import Data |
| 26 | + |
| 27 | + |
| 28 | +async def _store(session_id: str, flow_id, text: str) -> None: |
| 29 | + msg = Message(text=text, sender="User", sender_name="User", session_id=session_id) |
| 30 | + await aadd_messages([msg], flow_id=flow_id) |
| 31 | + |
| 32 | + |
| 33 | +async def test_aget_messages_defaults_flow_id_from_context(client): # noqa: ARG001 |
| 34 | + """The reproduction: two flows share a session_id; the ambient scope isolates them. |
| 35 | +
|
| 36 | + Simulates old frozen code calling ``aget_messages(session_id=...)`` with no flow_id while |
| 37 | + Flow B's graph is executing. The ambient flow scope must prevent Flow A's row from leaking. |
| 38 | + """ |
| 39 | + flow_a, flow_b = uuid4(), uuid4() |
| 40 | + session_id = "shared-session-13059" |
| 41 | + await _store(session_id, flow_a, "secret from A") |
| 42 | + await _store(session_id, flow_b, "hello from B") |
| 43 | + |
| 44 | + token = set_current_flow_id(flow_b) |
| 45 | + try: |
| 46 | + scoped = await aget_messages(session_id=session_id) |
| 47 | + finally: |
| 48 | + reset_current_flow_id(token) |
| 49 | + |
| 50 | + assert [m.text for m in scoped] == ["hello from B"] |
| 51 | + |
| 52 | + |
| 53 | +async def test_explicit_flow_id_overrides_context(client): # noqa: ARG001 |
| 54 | + """An explicitly-passed flow_id must win over the ambient context (no silent override).""" |
| 55 | + flow_a, flow_b = uuid4(), uuid4() |
| 56 | + session_id = "shared-session-13059-explicit" |
| 57 | + await _store(session_id, flow_a, "secret from A") |
| 58 | + await _store(session_id, flow_b, "hello from B") |
| 59 | + |
| 60 | + token = set_current_flow_id(flow_b) |
| 61 | + try: |
| 62 | + a_msgs = await aget_messages(session_id=session_id, flow_id=flow_a) |
| 63 | + finally: |
| 64 | + reset_current_flow_id(token) |
| 65 | + |
| 66 | + assert [m.text for m in a_msgs] == ["secret from A"] |
| 67 | + |
| 68 | + |
| 69 | +async def test_no_context_preserves_legacy_unscoped(client): # noqa: ARG001 |
| 70 | + """With no ambient scope and no explicit flow_id, behavior is unchanged (both rows).""" |
| 71 | + flow_a, flow_b = uuid4(), uuid4() |
| 72 | + session_id = "shared-session-13059-legacy" |
| 73 | + await _store(session_id, flow_a, "secret from A") |
| 74 | + await _store(session_id, flow_b, "hello from B") |
| 75 | + |
| 76 | + unscoped = await aget_messages(session_id=session_id) |
| 77 | + |
| 78 | + assert len(unscoped) == 2 |
| 79 | + |
| 80 | + |
| 81 | +async def test_context_accepts_string_flow_id(client): # noqa: ARG001 |
| 82 | + """``graph.flow_id`` is commonly a ``str``; the default path must coerce it, not crash. |
| 83 | +
|
| 84 | + ``MessageTable.flow_id`` is UUID-typed; on SQLite comparing it to a raw string makes the |
| 85 | + ``Uuid`` bind processor call ``value.hex`` and raise. The default must coerce str -> UUID. |
| 86 | + """ |
| 87 | + flow_a, flow_b = uuid4(), uuid4() |
| 88 | + session_id = "shared-session-13059-str" |
| 89 | + await _store(session_id, flow_a, "secret from A") |
| 90 | + await _store(session_id, flow_b, "hello from B") |
| 91 | + |
| 92 | + token = set_current_flow_id(str(flow_b)) |
| 93 | + try: |
| 94 | + scoped = await aget_messages(session_id=session_id) |
| 95 | + finally: |
| 96 | + reset_current_flow_id(token) |
| 97 | + |
| 98 | + assert [m.text for m in scoped] == ["hello from B"] |
| 99 | + |
| 100 | + |
| 101 | +async def test_invalid_context_flow_id_degrades_to_unscoped(client): # noqa: ARG001 |
| 102 | + """A non-UUID ambient flow_id (synthetic/test graphs) must degrade, never crash retrieval.""" |
| 103 | + flow_a, flow_b = uuid4(), uuid4() |
| 104 | + session_id = "shared-session-13059-badctx" |
| 105 | + await _store(session_id, flow_a, "secret from A") |
| 106 | + await _store(session_id, flow_b, "hello from B") |
| 107 | + |
| 108 | + token = set_current_flow_id("not-a-uuid") |
| 109 | + try: |
| 110 | + result = await aget_messages(session_id=session_id) |
| 111 | + finally: |
| 112 | + reset_current_flow_id(token) |
| 113 | + |
| 114 | + assert len(result) == 2 |
| 115 | + |
| 116 | + |
| 117 | +class _OldStyleMemory(MemoryComponent): |
| 118 | + """Simulates a flow saved before PR #13087: its frozen ``retrieve_messages`` omits ``flow_id``. |
| 119 | +
|
| 120 | + We cannot recreate frozen bytes in a unit test, so we reproduce the one behavior that matters: |
| 121 | + the internal-memory retrieve calls ``aget_messages`` with no ``flow_id``. The engine's ambient |
| 122 | + flow scope must still isolate it. |
| 123 | + """ |
| 124 | + |
| 125 | + name = "OldStyleMemory" |
| 126 | + |
| 127 | + async def retrieve_messages(self) -> Data: |
| 128 | + from langflow.memory import aget_messages as backend_aget_messages |
| 129 | + |
| 130 | + stored = await backend_aget_messages(session_id=self.session_id, order="DESC") |
| 131 | + return cast("Data", list(reversed(stored))) |
| 132 | + |
| 133 | + |
| 134 | +async def test_graph_execution_binds_flow_scope_end_to_end(client): # noqa: ARG001 |
| 135 | + """End-to-end: running Flow B's graph must not surface Flow A's message via unscoped frozen code. |
| 136 | +
|
| 137 | + This is the real reproduction path — ``get_instance_results`` binds ``graph.flow_id`` so the |
| 138 | + old-style component's unscoped ``aget_messages`` call is defaulted to Flow B's scope. |
| 139 | + """ |
| 140 | + flow_a, flow_b = uuid4(), uuid4() |
| 141 | + session_id = "shared-session-13059-e2e" |
| 142 | + await _store(session_id, flow_a, "SECRET_FROM_A") |
| 143 | + await _store(session_id, flow_b, "hello from B") |
| 144 | + |
| 145 | + probe = _OldStyleMemory(_id="old_memory") |
| 146 | + probe.set(session_id=session_id, n_messages=100, order="Ascending") |
| 147 | + chat_output = ChatOutput(_id="chat_output") |
| 148 | + chat_output.set(input_value=probe.retrieve_messages_as_text) |
| 149 | + |
| 150 | + graph = Graph(probe, chat_output, flow_id=str(flow_b)) |
| 151 | + async for _ in graph.async_start(): |
| 152 | + pass |
| 153 | + |
| 154 | + rendered = chat_output.get_output_by_method(chat_output.message_response).value |
| 155 | + text = rendered.text if hasattr(rendered, "text") else str(rendered) |
| 156 | + assert "SECRET_FROM_A" not in text, "Flow A's message leaked into Flow B's graph run" |
| 157 | + assert "hello from B" in text |
0 commit comments