fix: route memory ops to stubs when DB service is noop (#12808)

fix(lfx): route memory ops to stubs when DB service is noop

lfx.memory bound to langflow.memory at import time whenever the langflow
package was importable, even when the registered DB service was
NoopDatabaseService. langflow.memory.aupdate_messages then called
session.get() on a NoopSession (which unconditionally returns None) and
raised spurious "Message with id X not found" errors mid-stream from the
Agent component.

Dispatch now happens at call time via has_langflow_db_backend(), which
requires both langflow to be importable AND a non-noop DB service to be
registered. Call-time evaluation is required because the DB service is
typically registered after lfx.memory is first imported during component
class loading.

src/lfx/src/lfx/memory/__init__.py

@@ -1,60 +1,77 @@
-"""Memory management for lfx with dynamic loading.
+"""Memory management for lfx with dynamic dispatch.
 
-This module automatically chooses between the full langflow implementation
-(when available) and the lfx implementation (when standalone).
+Routes memory operations to either the full langflow implementation (when
+langflow is installed AND a real database service is registered) or the lfx
+stub implementation (standalone / noop DB).
+
+Dispatch is evaluated at call time, not import time, because the database
+service is typically registered *after* this module is first imported (e.g.,
+from Component class definitions loaded before graph setup). An import-time
+decision can't distinguish "langflow is importable" from "a real DB is wired",
+and picking the langflow backend with a NoopDatabaseService yields silent
+no-op inserts followed by spurious "Message with id X not found" errors on
+update.
 """
 
-from lfx.utils.langflow_utils import has_langflow_memory
+from __future__ import annotations
+
+from typing import Any
+
+from lfx.utils.langflow_utils import has_langflow_db_backend
+
+
+def _impl():
+    if has_langflow_db_backend():
+        from langflow import memory as impl
+    else:
+        from lfx.memory import stubs as impl
+    return impl
+
+
+def aadd_messages(*args: Any, **kwargs: Any):
+    return _impl().aadd_messages(*args, **kwargs)
+
+
+def aadd_messagetables(*args: Any, **kwargs: Any):
+    return _impl().aadd_messagetables(*args, **kwargs)
+
+
+def add_messages(*args: Any, **kwargs: Any):
+    return _impl().add_messages(*args, **kwargs)
+
+
+def adelete_messages(*args: Any, **kwargs: Any):
+    return _impl().adelete_messages(*args, **kwargs)
+
+
+def aget_messages(*args: Any, **kwargs: Any):
+    return _impl().aget_messages(*args, **kwargs)
+
+
+def astore_message(*args: Any, **kwargs: Any):
+    return _impl().astore_message(*args, **kwargs)
+
+
+def aupdate_messages(*args: Any, **kwargs: Any):
+    return _impl().aupdate_messages(*args, **kwargs)
+
+
+def delete_message(*args: Any, **kwargs: Any):
+    return _impl().delete_message(*args, **kwargs)
+
+
+def delete_messages(*args: Any, **kwargs: Any):
+    return _impl().delete_messages(*args, **kwargs)
+
+
+def get_messages(*args: Any, **kwargs: Any):
+    return _impl().get_messages(*args, **kwargs)
+
+
+def store_message(*args: Any, **kwargs: Any):
+    return _impl().store_message(*args, **kwargs)
 
-# Import the appropriate implementation
-if has_langflow_memory():
-    try:
-        # Import full langflow implementation
-        from langflow.memory import (
-            aadd_messages,
-            aadd_messagetables,
-            add_messages,
-            adelete_messages,
-            aget_messages,
-            astore_message,
-            aupdate_messages,
-            delete_message,
-            delete_messages,
-            get_messages,
-            store_message,
-        )
-    except ImportError:
-        # Fallback to lfx implementation if langflow import fails
-        from lfx.memory.stubs import (
-            aadd_messages,
-            aadd_messagetables,
-            add_messages,
-            adelete_messages,
-            aget_messages,
-            astore_message,
-            aupdate_messages,
-            delete_message,
-            delete_messages,
-            get_messages,
-            store_message,
-        )
-else:
-    # Use lfx implementation
-    from lfx.memory.stubs import (
-        aadd_messages,
-        aadd_messagetables,
-        add_messages,
-        adelete_messages,
-        aget_messages,
-        astore_message,
-        aupdate_messages,
-        delete_message,
-        delete_messages,
-        get_messages,
-        store_message,
-    )
 
-# Export the available functions
 __all__ = [
     "aadd_messages",
     "aadd_messagetables",

src/lfx/src/lfx/utils/langflow_utils.py

@@ -50,3 +50,22 @@ def has_langflow_memory():
     _LangflowModule.set_available(is_langflow_available)
 
     return is_langflow_available
+
+
+def has_langflow_db_backend() -> bool:
+    """Return True iff langflow-backed memory calls have a real DB to hit.
+
+    Requires both langflow to be importable AND the registered database
+    service to be a non-noop implementation. Evaluated on every call because
+    the database service is typically registered *after* this module is first
+    imported (e.g., from Component class definitions loaded before graph setup).
+    """
+    if not has_langflow_memory():
+        return False
+    from lfx.services.database.service import NoopDatabaseService
+    from lfx.services.deps import get_db_service
+
+    try:
+        return not isinstance(get_db_service(), NoopDatabaseService)
+    except Exception:  # noqa: BLE001
+        return False

src/lfx/tests/unit/memory/test_memory_dispatch.py

@@ -0,0 +1,118 @@
+"""Regression tests for lfx.memory runtime dispatch.
+
+Original bug: when langflow was installed alongside lfx but `lfx run` had
+only a NoopDatabaseService registered, `lfx.memory` bound at import time to
+`langflow.memory` (because the `langflow` package was importable). The
+langflow-backed `aupdate_messages` then called `session.get(...)` on a
+NoopSession, which always returns `None`, raising spurious
+"Message with id X not found" errors mid-stream.
+"""
+
+from __future__ import annotations
+
+import uuid
+
+import pytest
+from lfx.services.database.service import NoopDatabaseService
+from lfx.utils.langflow_utils import has_langflow_db_backend
+
+
+class _FakeRealDbService:
+    """Stand-in for any non-noop DatabaseService implementation."""
+
+
+class TestHasLangflowDbBackend:
+    def test_returns_false_when_langflow_not_importable(self, monkeypatch):
+        monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: False)
+        assert has_langflow_db_backend() is False
+
+    def test_returns_false_with_noop_db_service(self, monkeypatch):
+        monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: True)
+        monkeypatch.setattr("lfx.services.deps.get_db_service", lambda: NoopDatabaseService())
+        assert has_langflow_db_backend() is False
+
+    def test_returns_true_with_real_db_service(self, monkeypatch):
+        monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: True)
+        monkeypatch.setattr("lfx.services.deps.get_db_service", lambda: _FakeRealDbService())
+        assert has_langflow_db_backend() is True
+
+    def test_returns_false_when_get_db_service_raises(self, monkeypatch):
+        monkeypatch.setattr("lfx.utils.langflow_utils.has_langflow_memory", lambda: True)
+
+        def boom():
+            msg = "service manager exploded"
+            raise RuntimeError(msg)
+
+        monkeypatch.setattr("lfx.services.deps.get_db_service", boom)
+        assert has_langflow_db_backend() is False
+
+
+class TestMemoryDispatch:
+    def test_dispatches_to_stubs_when_no_real_db(self, monkeypatch):
+        import lfx.memory as memory_mod
+        from lfx.memory import stubs
+
+        monkeypatch.setattr("lfx.memory.has_langflow_db_backend", lambda: False)
+        assert memory_mod._impl() is stubs
+
+    def test_dispatches_to_langflow_when_real_db(self, monkeypatch):
+        pytest.importorskip("langflow.memory")
+        import langflow.memory as langflow_memory
+        import lfx.memory as memory_mod
+
+        monkeypatch.setattr("lfx.memory.has_langflow_db_backend", lambda: True)
+        assert memory_mod._impl() is langflow_memory
+
+    def test_dispatch_is_evaluated_per_call(self, monkeypatch):
+        """Dispatch must read the backend state each call, not cache at import.
+
+        The database service is often registered *after* lfx.memory is imported
+        (components load first, services register during graph setup), so
+        memoizing the dispatcher would bind to whatever state existed at
+        component-module load time.
+        """
+        import lfx.memory as memory_mod
+        from lfx.memory import stubs
+
+        state = {"real": False}
+        monkeypatch.setattr("lfx.memory.has_langflow_db_backend", lambda: state["real"])
+
+        assert memory_mod._impl() is stubs
+        state["real"] = True
+        pytest.importorskip("langflow.memory")
+        import langflow.memory as langflow_memory
+
+        assert memory_mod._impl() is langflow_memory
+
+
+class TestAupdateMessagesRegression:
+    """Direct regression for the original 'Message with id X not found' crash."""
+
+    @pytest.mark.asyncio
+    async def test_aupdate_messages_does_not_raise_against_noop_session(self, monkeypatch):
+        """Regression: route to stubs (no-op) instead of raising via langflow.memory.
+
+        With langflow importable but only a NoopDatabaseService registered,
+        aupdate_messages must route to stubs and succeed silently rather than
+        trigger langflow.memory's strict existence check against NoopSession.
+        """
+        try:
+            from langflow.schema.message import Message
+        except ImportError:
+            from lfx.schema.message import Message
+
+        # Force the noop-DB branch even if a real DB happens to be registered in
+        # this test environment.
+        monkeypatch.setattr("lfx.services.deps.get_db_service", lambda: NoopDatabaseService())
+
+        from lfx.memory import aupdate_messages
+
+        msg = Message(
+            id=str(uuid.uuid4()),
+            text="hello",
+            sender="AI",
+            sender_name="Test",
+            session_id="test-session",
+        )
+        result = await aupdate_messages(msg)
+        assert isinstance(result, list)