From f547adc69de85b551af93edb09b7f115e934dc70 Mon Sep 17 00:00:00 2001 From: Raffael Meyer <14891507+barredterra@users.noreply.github.com> Date: Mon, 13 Jul 2026 05:39:09 +0200 Subject: [PATCH] fix: isolate DB connection per tool call (#56) (cherry picked from commit 2e4a0b11349ac3a2ac71ff1e5fb16845464aec76) --- ask_alyf/ask_alyf/agent.py | 19 +++++---- ask_alyf/ask_alyf/test_code_tools.py | 44 +++++++++++++++++++- ask_alyf/ask_alyf/toolset.py | 61 +++++++++++++++++++++++++--- 3 files changed, 110 insertions(+), 14 deletions(-) diff --git a/ask_alyf/ask_alyf/agent.py b/ask_alyf/ask_alyf/agent.py index 39869d5..5011626 100644 --- a/ask_alyf/ask_alyf/agent.py +++ b/ask_alyf/ask_alyf/agent.py @@ -263,14 +263,17 @@ def _build_subagents(self) -> list[SubAgent]: ), "system_prompt": DOCUMENT_PLANNER_INSTRUCTIONS, "tools": [ - self.toolset.get_list, - self.toolset.get, - self.toolset.get_value, - self.toolset.get_single_value, - self.toolset.get_meta, - self.toolset.has_permission, - self.toolset.get_doc_permissions, - self.toolset.list_accessible_doctypes, + clear_messages_on_tool_error(fn) + for fn in ( + self.toolset.get_list, + self.toolset.get, + self.toolset.get_value, + self.toolset.get_single_value, + self.toolset.get_meta, + self.toolset.has_permission, + self.toolset.get_doc_permissions, + self.toolset.list_accessible_doctypes, + ) ], "response_format": DocumentPlannerResult, } diff --git a/ask_alyf/ask_alyf/test_code_tools.py b/ask_alyf/ask_alyf/test_code_tools.py index 41036cf..4f99227 100644 --- a/ask_alyf/ask_alyf/test_code_tools.py +++ b/ask_alyf/ask_alyf/test_code_tools.py @@ -1,6 +1,6 @@ import asyncio from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import frappe from frappe.tests import UnitTestCase @@ -556,6 +556,48 @@ def test_ask_mode_tools_exclude_all_write_proposals_and_host_mutation(self): self.assertFalse(proposal_tools.intersection(ask_tool_names)) self.assertFalse(host_mutation_tools.intersection(ask_tool_names)) + def test_clear_messages_wrapper_connects_without_changing_user(self): + inherited_db = frappe.local.db + private_db = MagicMock() + + def connect(*, set_admin_as_user): + self.assertFalse(set_admin_as_user) + frappe.local.db = private_db + + wrapped = clear_messages_on_tool_error(lambda: frappe.session.user) + with ( + patch("ask_alyf.ask_alyf.toolset.frappe.connect", side_effect=connect), + patch("ask_alyf.ask_alyf.toolset.frappe.set_user") as set_user, + ): + result = wrapped() + + self.assertEqual(result, frappe.session.user) + set_user.assert_not_called() + private_db.commit.assert_called_once_with() + private_db.close.assert_called_once_with() + self.assertIs(frappe.local.db, inherited_db) + + def test_clear_messages_wrapper_propagates_commit_failure(self): + inherited_db = frappe.local.db + private_db = MagicMock() + private_db.commit.side_effect = RuntimeError("commit failed") + + def connect(*, set_admin_as_user): + frappe.local.db = private_db + + wrapped = clear_messages_on_tool_error(lambda: "done") + with ( + patch("ask_alyf.ask_alyf.toolset.frappe.connect", side_effect=connect), + patch("ask_alyf.ask_alyf.toolset.frappe.clear_messages") as clear_messages, + self.assertRaisesRegex(RuntimeError, "commit failed"), + ): + wrapped() + + private_db.rollback.assert_called_once_with() + private_db.close.assert_called_once_with() + clear_messages.assert_called_once_with() + self.assertIs(frappe.local.db, inherited_db) + def test_clear_messages_wrapper_preserves_async_tools(self): async def fake_tool(file_id): return {"file_id": file_id} diff --git a/ask_alyf/ask_alyf/toolset.py b/ask_alyf/ask_alyf/toolset.py index fc658a8..6f93199 100644 --- a/ask_alyf/ask_alyf/toolset.py +++ b/ask_alyf/ask_alyf/toolset.py @@ -1,3 +1,4 @@ +import contextlib import functools import inspect from dataclasses import dataclass, field @@ -1011,17 +1012,62 @@ def show_chart( ) +@contextlib.contextmanager +def _isolated_db(): + """Run a block against a private Frappe DB connection. + + LangGraph runs sync tools via `asyncio.to_thread`, which copies the + calling context (contextvars) into the worker thread. Frappe's + `frappe.local` is contextvar-backed, so the worker thread inherits the + same `frappe.local.db` pymysql connection as the agent's main thread. + pymysql connections are not safe for concurrent use, so parallel tool + calls (and subagent tool calls) corrupt the shared connection's packet + stream — surfacing as `InterfaceError(0, '')` / `Packet sequence number + wrong`. + + This opens a fresh connection bound to the current thread/context, runs + the block, then closes it and restores the inherited binding. The agent + thread's own connection is never touched or closed. + + No-op when Frappe isn't initialized (e.g. direct unit-test calls). + """ + conf = getattr(frappe.local, "conf", None) + if not conf or not getattr(conf, "db_name", None): + yield + return + + inherited_db = getattr(frappe.local, "db", None) + + frappe.connect(set_admin_as_user=False) + try: + yield + frappe.db.commit() + except Exception: + with contextlib.suppress(Exception): + frappe.db.rollback() + raise + finally: + with contextlib.suppress(Exception): + frappe.local.db.close() + frappe.local.db = inherited_db + + def clear_messages_on_tool_error(func): - """Wrap a tool callable so that queued Frappe messages are discarded on - exception. The error still propagates to the agent framework (so the LLM - sees it), but the user won't receive a popup.""" + """Wrap a tool so each call runs on a private DB connection and queued + Frappe messages are discarded on exception. + + The private connection (see `_isolated_db`) is what makes tool calls + safe under LangGraph's threaded tool executor; the error handling keeps + user-facing popups from firing when a tool fails inside the agent loop. + """ if inspect.iscoroutinefunction(func): @functools.wraps(func) async def async_wrapper(*args, **kwargs): try: - return await func(*args, **kwargs) + with _isolated_db(): + return await func(*args, **kwargs) except Exception: frappe.clear_messages() raise @@ -1031,9 +1077,14 @@ async def async_wrapper(*args, **kwargs): @functools.wraps(func) def wrapper(*args, **kwargs): try: - return func(*args, **kwargs) + return _run_with_isolated_db(func, args, kwargs) except Exception: frappe.clear_messages() raise return wrapper + + +def _run_with_isolated_db(func, args, kwargs): + with _isolated_db(): + return func(*args, **kwargs)