From 6f3e6c1f9e178abf6b904cf13c58120818be25ed Mon Sep 17 00:00:00 2001 From: Shubham Padkonde Date: Fri, 6 Nov 2026 01:02:29 +0530 Subject: [PATCH 1/2] fix: refresh AI tool availability after provider preference changes Update the shared Query Tool flag and existing report menus from the latest status response when the default provider changes. Cover provider removal, repeated saves, and out-of-order responses. --- web/pgadmin/llm/static/js/ai_tools.js | 53 +++++++----- web/pgadmin/tools/sqleditor/__init__.py | 3 + .../js/components/QueryToolComponent.jsx | 3 +- .../tests/test_panel_ai_availability.py | 35 ++++++++ .../javascript/llm/ai_tools_spec.js | 83 +++++++++++++++++++ 5 files changed, 154 insertions(+), 23 deletions(-) create mode 100644 web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py create mode 100644 web/regression/javascript/llm/ai_tools_spec.js diff --git a/web/pgadmin/llm/static/js/ai_tools.js b/web/pgadmin/llm/static/js/ai_tools.js index b259fb5b202..9dc73c205b4 100644 --- a/web/pgadmin/llm/static/js/ai_tools.js +++ b/web/pgadmin/llm/static/js/ai_tools.js @@ -11,6 +11,8 @@ import AIReport from './AIReport'; import { AllPermissionTypes, BROWSER_PANELS } from '../../../browser/static/js/constants'; import getApiInstance from '../../../static/js/api_instance'; import MainMenuFactory from '../../../browser/static/js/MainMenuFactory'; +import usePreferences from '../../../preferences/static/js/store'; +import pgAdmin from 'sources/pgadmin'; import url_for from 'sources/url_for'; // AI Reports Module @@ -37,7 +39,15 @@ define([ this.initialized = true; - // Check LLM status and only register menus if enabled + const getProvider = (state) => state.getPreferences('ai', 'default_provider')?.value; + let provider = getProvider(usePreferences.getState()); + usePreferences.subscribe((state) => { + const nextProvider = getProvider(state); + if (provider !== nextProvider) { + provider = nextProvider; + this.checkLLMStatus(); + } + }); this.checkLLMStatus(); return this; @@ -167,27 +177,26 @@ define([ }, // Check if LLM is configured, register menus only if system-enabled - checkLLMStatus: function() { - const api = getApiInstance(); - api.get(url_for('llm.status')) - .then((res) => { - if (res?.data?.success) { - this.llmEnabled = res.data.data?.enabled || false; - this.llmSystemEnabled = res.data.data?.system_enabled || false; - } - this.llmStatusChecked = true; - - // Only register menus if LLM is enabled at system level - if (this.llmSystemEnabled) { - this.registerMenus(); - MainMenuFactory.createMainMenus(); - } - }) - .catch(() => { - this.llmEnabled = false; - this.llmSystemEnabled = false; - this.llmStatusChecked = true; - }); + checkLLMStatus: async function() { + const request = getApiInstance().get(url_for('llm.status')); + this.statusRequest = request; + const res = await request.catch(() => null); + // An earlier save can finish after a newer one. + if (request !== this.statusRequest) return; + this.statusRequest = null; + const status = res?.data?.success ? res.data.data : null; + this.llmEnabled = Boolean(status?.enabled); + this.llmSystemEnabled = Boolean(status?.system_enabled); + this.llmStatusChecked = true; + pgAdmin.llm_enabled = this.llmEnabled; + + if (this.llmSystemEnabled && !this.menusRegistered) { + this.registerMenus(); + this.menusRegistered = true; + MainMenuFactory.createMainMenus(); + } else if (this.menusRegistered) { + MainMenuFactory.enableDisableMenus(); + } }, // Get the node type from tree item diff --git a/web/pgadmin/tools/sqleditor/__init__.py b/web/pgadmin/tools/sqleditor/__init__.py index 8fa4c9204e2..51f3d557325 100644 --- a/web/pgadmin/tools/sqleditor/__init__.py +++ b/web/pgadmin/tools/sqleditor/__init__.py @@ -398,6 +398,8 @@ def panel(trans_id): trans_id: unique transaction id """ + from pgadmin.llm.utils import is_llm_enabled + params = None if request.args: params = {k: v for k, v in request.args.items()} @@ -442,6 +444,7 @@ def panel(trans_id): params['user'] = underscore_escape(s.username) if 'role' not in params and s.role: params['role'] = underscore_escape(s.role) + params['llm_enabled'] = is_llm_enabled() params['layout'] = get_setting('SQLEditor/Layout') params['macros'] = get_user_macros() params['is_desktop_mode'] = current_app.PGADMIN_RUNTIME diff --git a/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx b/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx index 49cb7f5ffc6..80f152e5bed 100644 --- a/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx +++ b/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx @@ -258,7 +258,7 @@ export default function QueryToolComponent({params, pgWindow, pgAdmin, selectedN maximizable: true, tabs: [ LayoutDocker.getPanel({id: PANELS.QUERY, title: gettext('Query'), content: setSelectedText(text)} setQtStatePartial={setQtStatePartial}/>}), - ...(pgAdmin.llm_enabled ? [LayoutDocker.getPanel({id: PANELS.AI_ASSISTANT, title: gettext('AI Assistant'), content: })] : []), + ...(params.llm_enabled ? [LayoutDocker.getPanel({id: PANELS.AI_ASSISTANT, title: gettext('AI Assistant'), content: })] : []), LayoutDocker.getPanel({id: PANELS.HISTORY, title: gettext('Query History'), content: }), ], }, @@ -1159,6 +1159,7 @@ QueryToolComponent.propTypes = { server_name: PropTypes.string, database_name: PropTypes.string, layout: PropTypes.string, + llm_enabled: PropTypes.bool, fileName: PropTypes.string, storage: PropTypes.string, }), diff --git a/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py b/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py new file mode 100644 index 00000000000..46f8ab6c9d4 --- /dev/null +++ b/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py @@ -0,0 +1,35 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +import json +from types import SimpleNamespace +from unittest.mock import patch + +from pgadmin.utils.route import BaseTestGenerator + + +class TestPanelAIAvailability(BaseTestGenerator): + """Each new Query Tool layout gets the current server-side AI status.""" + + @patch('pgadmin.tools.sqleditor.render_template', return_value='') + @patch('pgadmin.tools.sqleditor.get_server') + @patch('pgadmin.llm.utils.is_llm_enabled') + def runTest(self, enabled, get_server, render): + get_server.return_value = SimpleNamespace( + shared=False, bgcolor=None, name='test', username='test', role='') + for value in (True, False): + with self.subTest(enabled=value): + enabled.return_value = value + response = self.tester.post('/sqleditor/panel/1', query_string={ + 'sid': 1, 'title': 'Query Tool', 'selectedNodeInfo': '{}', + 'llm_enabled': str(not value).lower(), + }) + self.assertEqual(response.status_code, 200) + params = json.loads(render.call_args.kwargs['params']) + self.assertIs(params['llm_enabled'], value) diff --git a/web/regression/javascript/llm/ai_tools_spec.js b/web/regression/javascript/llm/ai_tools_spec.js new file mode 100644 index 00000000000..e39e0e99a7a --- /dev/null +++ b/web/regression/javascript/llm/ai_tools_spec.js @@ -0,0 +1,83 @@ +///////////////////////////////////////////////////////////// +// +// pgAdmin 4 - PostgreSQL Tools +// +// Copyright (C) 2013 - 2026, The pgAdmin Development Team +// This software is released under the PostgreSQL Licence +// +////////////////////////////////////////////////////////////// + +import {waitFor} from '@testing-library/react'; + +jest.unmock('zustand'); +jest.mock('../../../pgadmin/llm/static/js/AIReport', () => jest.fn()); +jest.mock('../../../pgadmin/static/js/api_instance', () => () => ({get: mockGet})); +jest.mock('sources/url_for', () => (name) => name); +jest.mock('sources/pgadmin', () => ({})); +jest.mock('../../../pgadmin/browser/static/js/MainMenuFactory', () => ({ + createMainMenus: jest.fn(), enableDisableMenus: jest.fn(), +})); + +const mockGet = jest.fn(); +const status = (enabled) => ({data: {success: true, data: {enabled, system_enabled: true}}}); + +describe('AI tools after preferences are saved', () => { + let browser, pgAdmin, preferences, tools; + + beforeEach(() => { + jest.resetModules(); + mockGet.mockReset(); + browser = {add_menus: jest.fn(), add_menu_category: jest.fn()}; + pgAdmin = require('sources/pgadmin'); + preferences = require('../../../pgadmin/preferences/static/js/store').default; + global.define = (_deps, factory) => factory((s) => s, browser); + try { + require('../../../pgadmin/llm/static/js/ai_tools'); + } finally { + delete global.define; + } + tools = browser.AITools; + }); + + // PreferencesComponent calls cache() after a successful save. + async function refreshPreferences(provider) { + mockGet.mockResolvedValueOnce({data: [{module: 'ai', name: 'default_provider', value: provider}]}); + await preferences.getState().cache(); + } + + it('updates report menus and the flag for new Query Tools without registering duplicate menus', async () => { + await refreshPreferences(''); + mockGet.mockResolvedValue(status(false)); + tools.init(); + await waitFor(() => expect(tools.llmStatusChecked).toBe(true)); + expect(tools.checkLLMEnabled()).toBe(false); + mockGet.mockResolvedValue(status(true)); + await refreshPreferences('ollama'); + await waitFor(() => expect(pgAdmin.llm_enabled).toBe(true)); + expect(tools.checkLLMEnabled()).toBe(true); + + tools.init(); + const calls = mockGet.mock.calls.length; + await refreshPreferences('ollama'); + expect(mockGet.mock.calls.slice(calls)).toEqual([['preferences.get_all']]); + mockGet.mockResolvedValue(status(false)); + await refreshPreferences(''); + await waitFor(() => expect(pgAdmin.llm_enabled).toBe(false)); + expect(tools.checkLLMEnabled()).toBe(false); + expect(browser.add_menus).toHaveBeenCalledTimes(1); + }); + + it('keeps a failed lookup disabled even if an older successful request finishes later', async () => { + let resolveFirst; + mockGet.mockReturnValueOnce(new Promise(resolve => { resolveFirst = resolve; })); + tools.init(); + mockGet.mockRejectedValue(new Error('status unavailable')); + await refreshPreferences('openai'); + await waitFor(() => expect(tools.llmStatusChecked).toBe(true)); + resolveFirst(status(true)); + await new Promise(resolve => setTimeout(resolve, 0)); + expect(tools.checkLLMEnabled()).toBe(false); + expect(pgAdmin.llm_enabled).toBe(false); + expect(browser.add_menus).not.toHaveBeenCalled(); + }); +}); From 42cc6763b5e8cc4bd9e3f48395ac52693371b4c1 Mon Sep 17 00:00:00 2001 From: Dave Page Date: Fri, 6 Nov 2026 16:43:32 +0000 Subject: [PATCH 2/2] Wrap an overlong line in the AI panel availability test to satisfy pycodestyle. --- .../sqleditor/tests/test_panel_ai_availability.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py b/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py index 46f8ab6c9d4..e4524c422c2 100644 --- a/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py +++ b/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py @@ -26,10 +26,13 @@ def runTest(self, enabled, get_server, render): for value in (True, False): with self.subTest(enabled=value): enabled.return_value = value - response = self.tester.post('/sqleditor/panel/1', query_string={ - 'sid': 1, 'title': 'Query Tool', 'selectedNodeInfo': '{}', - 'llm_enabled': str(not value).lower(), - }) + response = self.tester.post( + '/sqleditor/panel/1', + query_string={ + 'sid': 1, 'title': 'Query Tool', + 'selectedNodeInfo': '{}', + 'llm_enabled': str(not value).lower(), + }) self.assertEqual(response.status_code, 200) params = json.loads(render.call_args.kwargs['params']) self.assertIs(params['llm_enabled'], value)