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..e4524c422c2 --- /dev/null +++ b/web/pgadmin/tools/sqleditor/tests/test_panel_ai_availability.py @@ -0,0 +1,38 @@ +########################################################################## +# +# 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(); + }); +});