diff --git a/src/backend/base/langflow/api/v1/mcp_utils.py b/src/backend/base/langflow/api/v1/mcp_utils.py index 5407e1fbb0..f01e0f57cc 100644 --- a/src/backend/base/langflow/api/v1/mcp_utils.py +++ b/src/backend/base/langflow/api/v1/mcp_utils.py @@ -296,10 +296,8 @@ async def handle_call_tool( progress_token=progress_token, progress=0.0, total=1.0 ) - conversation_id = str(uuid4()) - input_request = SimplifiedAPIRequest( - input_value=processed_inputs.get("input_value", ""), session_id=conversation_id - ) + session_id = processed_inputs.pop("session_id", None) or str(uuid4()) + input_request = SimplifiedAPIRequest(input_value=processed_inputs.get("input_value", ""), session_id=session_id) async def send_progress_updates(progress_token): try: diff --git a/src/backend/base/langflow/api/v1/memories.py b/src/backend/base/langflow/api/v1/memories.py index 5a88ea00df..2d79385763 100644 --- a/src/backend/base/langflow/api/v1/memories.py +++ b/src/backend/base/langflow/api/v1/memories.py @@ -5,7 +5,7 @@ Endpoints: GET /memories - List (current user, paginated) GET /memories/{id} - Get one GET /memories/{id}/sessions - List sessions (tracked + untracked from MessageTable) - PATCH /memories/{id} - Update (name / threshold / auto_capture / preprocessing) + PATCH /memories/{id} - Update (name / threshold / auto_capture) DELETE /memories/{id} - Delete (cancels active tasks + removes KB from disk) POST /memories/{id}/flush - Manual flush / trigger ingestion POST /memories/{id}/regenerate - Regenerate from mismatch @@ -257,15 +257,17 @@ async def update_memory_base( current_user: CurrentActiveUser, patch: Annotated[MemoryBaseUpdate, Body(embed=False)] = ..., ) -> MemoryBaseRead: - """Update mutable parameters (threshold, auto_capture, preprocessing, etc.). + """Update mutable parameters (name, threshold, auto_capture). Threshold changes only take effect at the next auto-capture trigger. Any already-running ingestion task continues with its original arguments. + Preprocessing fields (preprocessing, preproc_model, preproc_instructions, preproc_kill_phrase) + are immutable after creation and cannot be patched. """ try: mb = await get_memory_base_service().update(memory_base_id, user_id=current_user.id, patch=patch) except PreprocessingValidationError as exc: - raise HTTPException(status_code=422, detail=str(exc)) from exc + raise HTTPException(status_code=403, detail=str(exc)) from exc if mb is None: raise HTTPException(status_code=404, detail="Memory base not found") return MemoryBaseRead.model_validate(mb) diff --git a/src/backend/base/langflow/helpers/flow.py b/src/backend/base/langflow/helpers/flow.py index bf98db5d05..a175c8e6a3 100644 --- a/src/backend/base/langflow/helpers/flow.py +++ b/src/backend/base/langflow/helpers/flow.py @@ -500,4 +500,13 @@ def json_schema_from_flow(flow: Flow) -> dict: if field_data.get("required", False): required.append(field_name) + if "session_id" not in properties: + properties["session_id"] = { + "type": "string", + "description": ( + "Optional session identifier used to persist conversation " + "history across tool calls. Omit to start a new session." + ), + } + return {"type": "object", "properties": properties, "required": required} diff --git a/src/backend/base/langflow/services/database/models/memory_base/model.py b/src/backend/base/langflow/services/database/models/memory_base/model.py index 244c511d0a..b8305cd99f 100644 --- a/src/backend/base/langflow/services/database/models/memory_base/model.py +++ b/src/backend/base/langflow/services/database/models/memory_base/model.py @@ -60,10 +60,6 @@ class MemoryBaseUpdate(SQLModel): name: str | None = None threshold: int | None = None auto_capture: bool | None = None - preprocessing: bool | None = None - preproc_model: str | None = None - preproc_instructions: str | None = None - preproc_kill_phrase: str | None = None class MemoryBaseRead(MemoryBaseBase): diff --git a/src/backend/base/langflow/services/memory_base/service.py b/src/backend/base/langflow/services/memory_base/service.py index d62b2f9689..15c1b5bd65 100644 --- a/src/backend/base/langflow/services/memory_base/service.py +++ b/src/backend/base/langflow/services/memory_base/service.py @@ -184,11 +184,8 @@ class MemoryBaseService(Service): if mb is None: return None - # Validate API key before committing if preprocessing will be active after the patch. - effective_preprocessing = patch.preprocessing if patch.preprocessing is not None else mb.preprocessing - effective_preproc_model = patch.preproc_model if patch.preproc_model is not None else mb.preproc_model - if effective_preprocessing: - _validate_preprocessing_api_key(user_id, effective_preproc_model) + if mb.preprocessing: + _validate_preprocessing_api_key(user_id, mb.preproc_model) for field, value in patch.model_dump(exclude_unset=True).items(): setattr(mb, field, value) diff --git a/src/backend/tests/unit/api/v1/test_mcp_utils.py b/src/backend/tests/unit/api/v1/test_mcp_utils.py index 1374b08b56..602c46e45b 100644 --- a/src/backend/tests/unit/api/v1/test_mcp_utils.py +++ b/src/backend/tests/unit/api/v1/test_mcp_utils.py @@ -1,8 +1,10 @@ from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest from langflow.api.utils.core import extract_global_variables_from_headers from langflow.api.v1 import mcp_utils +from langflow.helpers import flow as flow_helpers from lfx.interface.components import component_cache @@ -370,3 +372,160 @@ async def test_handle_list_tools_requires_current_user_on_global_server(monkeypa # No user context set — must return empty. tools = await mcp_utils.handle_list_tools() assert tools == [] + + +# ============================================================================ +# session_id propagation — MCP clients must be able to persist chat history. +# ============================================================================ + + +def _build_fake_server() -> SimpleNamespace: + """Build a minimal MCP server stub with progress notifications disabled.""" + return SimpleNamespace( + request_context=SimpleNamespace( + meta=SimpleNamespace(progressToken=None), + session=SimpleNamespace(send_progress_notification=AsyncMock()), + ) + ) + + +async def _invoke_handle_call_tool(monkeypatch, arguments: dict) -> AsyncMock: + """Run handle_call_tool with all external deps stubbed; return the simple_run_flow mock.""" + flow = SimpleNamespace(id="flow-id-1", name="my_flow", folder_id=None) + + async def fake_get_flow_snake_case(*_args, **_kwargs): + return flow + + run_response = SimpleNamespace(outputs=[]) + simple_run_flow_mock = AsyncMock(return_value=run_response) + + monkeypatch.setattr(mcp_utils, "get_flow_snake_case", fake_get_flow_snake_case) + monkeypatch.setattr(mcp_utils, "simple_run_flow", simple_run_flow_mock) + monkeypatch.setattr(mcp_utils, "with_db_session", lambda operation: operation(SimpleNamespace())) + # Force progress notifications off so the test does not exercise that path. + mcp_utils.get_mcp_config().enable_progress_notifications = False + + token = mcp_utils.current_user_ctx.set(SimpleNamespace(id="user-1")) + try: + await mcp_utils.handle_call_tool( + name="my_flow", + arguments=arguments, + server=_build_fake_server(), + ) + finally: + mcp_utils.current_user_ctx.reset(token) + + return simple_run_flow_mock + + +@pytest.mark.asyncio +async def test_handle_call_tool_uses_provided_session_id(monkeypatch): + """When the MCP client supplies session_id, it must be forwarded to simple_run_flow.""" + simple_run_flow_mock = await _invoke_handle_call_tool( + monkeypatch, + arguments={"input_value": "hello", "session_id": "user-1-thread-7"}, + ) + + simple_run_flow_mock.assert_awaited_once() + forwarded_request = simple_run_flow_mock.await_args.kwargs["input_request"] + assert forwarded_request.session_id == "user-1-thread-7" + assert forwarded_request.input_value == "hello" + + +@pytest.mark.asyncio +async def test_handle_call_tool_generates_session_id_when_omitted(monkeypatch): + """When session_id is absent or blank, a non-empty fallback id must be generated.""" + from uuid import UUID + + simple_run_flow_mock = await _invoke_handle_call_tool( + monkeypatch, + arguments={"input_value": "hello"}, + ) + + forwarded_request = simple_run_flow_mock.await_args.kwargs["input_request"] + # Fallback must be a valid UUID-shaped string. + UUID(forwarded_request.session_id) + + +@pytest.mark.asyncio +async def test_handle_call_tool_falls_back_when_session_id_blank(monkeypatch): + """Empty-string session_id must trigger the UUID fallback, not pass through.""" + from uuid import UUID + + simple_run_flow_mock = await _invoke_handle_call_tool( + monkeypatch, + arguments={"input_value": "hello", "session_id": ""}, + ) + + forwarded_request = simple_run_flow_mock.await_args.kwargs["input_request"] + UUID(forwarded_request.session_id) + assert forwarded_request.session_id != "" + + +def test_json_schema_from_flow_includes_optional_session_id(monkeypatch): + """json_schema_from_flow must advertise session_id so MCP clients can supply it.""" + + class _FakeGraph: + def __init__(self): + self.vertices = [] # No input nodes — exercises the empty-properties path. + + @classmethod + def from_payload(cls, _flow_data): + return cls() + + # Patch the lazy import inside json_schema_from_flow. + import lfx.graph.graph.base as graph_base_module + + monkeypatch.setattr(graph_base_module, "Graph", _FakeGraph) + + flow = SimpleNamespace(data={"nodes": [], "edges": []}) + schema = flow_helpers.json_schema_from_flow(flow) + + assert schema["type"] == "object" + assert "session_id" in schema["properties"] + assert schema["properties"]["session_id"]["type"] == "string" + # session_id must be optional so existing MCP clients keep working. + assert "session_id" not in schema["required"] + + +def test_json_schema_from_flow_preserves_flow_defined_session_id(monkeypatch): + """If a flow already defines a session_id input, do not overwrite it.""" + custom_session_id_property = { + "type": "string", + "description": "Flow-defined session id with custom semantics.", + } + + class _FakeNode: + is_input = True + data = { + "node": { + "template": { + "session_id": { + "show": True, + "advanced": False, + "type": "str", + "info": custom_session_id_property["description"], + "required": True, + } + } + } + } + + class _FakeGraph: + def __init__(self): + self.vertices = [_FakeNode()] + + @classmethod + def from_payload(cls, _flow_data): + return cls() + + import lfx.graph.graph.base as graph_base_module + + monkeypatch.setattr(graph_base_module, "Graph", _FakeGraph) + + flow = SimpleNamespace(data={"nodes": [], "edges": []}) + schema = flow_helpers.json_schema_from_flow(flow) + + # The flow's own definition wins — the reserved injection must not clobber it. + assert schema["properties"]["session_id"]["description"] == custom_session_id_property["description"] + assert "session_id" in schema["required"] diff --git a/src/backend/tests/unit/test_memory_bases.py b/src/backend/tests/unit/test_memory_bases.py index b840500e68..d1e7c69232 100644 --- a/src/backend/tests/unit/test_memory_bases.py +++ b/src/backend/tests/unit/test_memory_bases.py @@ -1601,7 +1601,7 @@ class TestMemoriesAPIHandlers: assert "API key" in exc_info.value.detail @pytest.mark.asyncio - async def test_update_missing_api_key_returns_422(self, mock_user): + async def test_update_missing_api_key_returns_403(self, mock_user): from fastapi import HTTPException from langflow.api.v1.memories import update_memory_base from langflow.services.memory_base.service import PreprocessingValidationError @@ -1615,10 +1615,10 @@ class TestMemoriesAPIHandlers: await update_memory_base( memory_base_id=uuid.uuid4(), current_user=mock_user, - patch=MemoryBaseUpdate(preprocessing=True, preproc_model="gpt-4o"), + patch=MemoryBaseUpdate(threshold=10), ) - assert exc_info.value.status_code == 422 + assert exc_info.value.status_code == 403 assert "API key" in exc_info.value.detail # ---------------------------------------------------------------- # @@ -2073,7 +2073,10 @@ class TestPreprocessingApiKeyValidation: @pytest.mark.asyncio async def test_service_update_raises_preprocessing_validation_error_on_missing_key(self): - """update() raises PreprocessingValidationError when enabling preprocessing without a key.""" + """update() raises PreprocessingValidationError. + + When the existing MB uses preprocessing but the API key is gone. + """ from langflow.services.memory_base.service import ( MemoryBaseService, PreprocessingValidationError, @@ -2082,8 +2085,8 @@ class TestPreprocessingApiKeyValidation: service = MemoryBaseService() user_id = uuid.uuid4() mb = _make_mb(user_id=user_id) - mb.preprocessing = False - mb.preproc_model = None + mb.preprocessing = True + mb.preproc_model = "gpt-4o" exec_result = MagicMock() exec_result.first.return_value = mb @@ -2114,7 +2117,7 @@ class TestPreprocessingApiKeyValidation: await service.update( mb.id, user_id, - MemoryBaseUpdate(preprocessing=True, preproc_model="gpt-4o"), + MemoryBaseUpdate(threshold=10), )