fix(tools): use short-lived sessions for icon lookups to prevent idle-in-transaction (#36903)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
goingforstudying-ctrl
2026-06-01 21:59:10 -04:00
committed by GitHub
parent a3265f722e
commit c64d3e98c4
2 changed files with 51 additions and 25 deletions
+25 -18
View File
@@ -962,34 +962,41 @@ class ToolManager:
@classmethod
def generate_workflow_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict:
try:
workflow_provider: WorkflowToolProvider | None = db.session.scalar(
select(WorkflowToolProvider)
.where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id)
.limit(1)
)
# Use a short-lived session to avoid holding a database transaction
# during long-running nested workflow execution.
# Fixes: idle in transaction when Workflow Tool runs (#36902)
with Session(db.engine, expire_on_commit=False) as session:
workflow_provider: WorkflowToolProvider | None = session.scalar(
select(WorkflowToolProvider)
.where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id)
.limit(1)
)
if workflow_provider is None:
raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
if workflow_provider is None:
raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
icon = emoji_icon_adapter.validate_json(workflow_provider.icon)
return icon
icon = emoji_icon_adapter.validate_json(workflow_provider.icon)
return icon
except Exception:
return {"background": "#252525", "content": "\ud83d\ude01"}
@classmethod
def generate_api_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict:
try:
api_provider: ApiToolProvider | None = db.session.scalar(
select(ApiToolProvider)
.where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id)
.limit(1)
)
# Use a short-lived session to avoid holding a database transaction
# during long-running tool execution.
with Session(db.engine, expire_on_commit=False) as session:
api_provider: ApiToolProvider | None = session.scalar(
select(ApiToolProvider)
.where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id)
.limit(1)
)
if api_provider is None:
raise ToolProviderNotFoundError(f"api provider {provider_id} not found")
if api_provider is None:
raise ToolProviderNotFoundError(f"api provider {provider_id} not found")
icon = emoji_icon_adapter.validate_json(api_provider.icon)
return icon
icon = emoji_icon_adapter.validate_json(api_provider.icon)
return icon
except Exception:
return {"background": "#252525", "content": "\ud83d\ude01"}
@@ -6,7 +6,7 @@ import json
import threading
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock, patch
from unittest.mock import MagicMock, Mock, patch
import pytest
@@ -802,17 +802,36 @@ def test_generate_tool_icon_urls_for_builtin_and_plugin():
def test_generate_tool_icon_urls_for_workflow_and_api():
workflow_provider = SimpleNamespace(icon='{"background": "#222", "content": "W"}')
api_provider = SimpleNamespace(icon='{"background": "#333", "content": "A"}')
mock_engine = object()
with patch("core.tools.tool_manager.db") as mock_db:
mock_db.session.scalar.side_effect = [workflow_provider, api_provider]
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "wf-1") == {"background": "#222", "content": "W"}
assert ToolManager.generate_api_tool_icon_url("tenant-1", "api-1") == {"background": "#333", "content": "A"}
mock_db.engine = mock_engine
with patch("core.tools.tool_manager.Session") as mock_session_cls:
mock_session = MagicMock()
mock_session.scalar.side_effect = [workflow_provider, api_provider]
mock_session_cls.return_value.__enter__ = MagicMock(return_value=mock_session)
mock_session_cls.return_value.__exit__ = MagicMock(return_value=False)
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "wf-1") == {
"background": "#222",
"content": "W",
}
assert ToolManager.generate_api_tool_icon_url("tenant-1", "api-1") == {"background": "#333", "content": "A"}
# Verify sessions are created with the engine
assert mock_session_cls.call_count == 2
mock_session_cls.assert_called_with(mock_engine, expire_on_commit=False)
def test_generate_tool_icon_urls_missing_workflow_and_api_use_default():
mock_engine = object()
with patch("core.tools.tool_manager.db") as mock_db:
mock_db.session.scalar.return_value = None
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
assert ToolManager.generate_api_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
mock_db.engine = mock_engine
with patch("core.tools.tool_manager.Session") as mock_session_cls:
mock_session = MagicMock()
mock_session.scalar.return_value = None
mock_session_cls.return_value.__enter__ = MagicMock(return_value=mock_session)
mock_session_cls.return_value.__exit__ = MagicMock(return_value=False)
assert ToolManager.generate_workflow_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
assert ToolManager.generate_api_tool_icon_url("tenant-1", "missing")["background"] == "#252525"
assert mock_session_cls.call_count == 2
def test_get_tool_icon_for_builtin_provider_variants():