feat(query row): add query with Assistant action for Loki and Prometheus (#116275)

* feat(query): query with Assistant action

* Restyle and add feature toggle

* Add test

* Update

* use module mapper

* Translations

* Use feature flag config

* Improve naming of props

* Remove queryWithAssistant feature toggle

* Remove feature toggle

* Restore queryWithAssistant feature toggle

* Add Expression property to queryWithAssistant feature toggle

* Bring back feature toggle

* lint

jest.config.js

@@ -73,6 +73,8 @@ module.exports = {
     // prevent systemjs amd extra from breaking tests.
     'systemjs/dist/extras/amd': '<rootDir>/public/test/mocks/systemjsAMDExtra.ts',
     '@bsull/augurs': '<rootDir>/public/test/mocks/augurs.ts',
+    // Mock @grafana/assistant to prevent initialization errors in tests
+    '^@grafana/assistant$': '<rootDir>/public/test/mocks/assistant.ts',
   },
   // Log the test results with dynamic Loki tags. Drone CI only
   reporters: ['default', ['<rootDir>/public/test/log-reporter.js', { enable: process.env.DRONE === 'true' }]],

packages/grafana-data/src/types/featureToggles.gen.ts

@@ -1420,4 +1420,9 @@ export interface FeatureToggles {
   * @default false
   */
   alertingSyncDispatchTimer?: boolean;
+  /**
+  * Enables the Query with Assistant button in the query editor
+  * @default false
+  */
+  queryWithAssistant?: boolean;
 }

pkg/services/featuremgmt/registry.go

@@ -2240,6 +2240,14 @@ var (
 			HideFromDocs:    true,
 			Expression:      "false",
 		},
+		{
+			Name:         "queryWithAssistant",
+			Description:  "Enables the Query with Assistant button in the query editor",
+			Stage:        FeatureStageExperimental,
+			FrontendOnly: true,
+			Owner:        grafanaOSSBigTent,
+			Expression:   "false",
+		},
 	}
 )
 

pkg/services/featuremgmt/toggles_gen.csv

@@ -46,7 +46,7 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2023-09-28,externalServiceAccounts,preview,@grafana/identity-access-team,false,false,false
 2023-10-03,enableNativeHTTPHistogram,experimental,@grafana/grafana-backend-services-squad,false,true,false
 2024-06-18,disableClassicHTTPHistogram,experimental,@grafana/grafana-backend-services-squad,false,true,false
-2023-12-06,kubernetesSnapshots,experimental,@grafana/grafana-app-platform-squad,false,true,false
+2023-12-05,kubernetesSnapshots,experimental,@grafana/grafana-app-platform-squad,false,true,false
 2025-06-26,kubernetesLibraryPanels,experimental,@grafana/grafana-app-platform-squad,false,true,false
 2024-06-05,kubernetesDashboards,GA,@grafana/dashboards-squad,false,false,false
 2025-08-01,kubernetesShortURLs,experimental,@grafana/grafana-app-platform-squad,false,true,false
@@ -54,7 +54,7 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2025-08-01,kubernetesAlertingRules,experimental,@grafana/alerting-squad,false,true,false
 2025-08-29,kubernetesCorrelations,experimental,@grafana/datapro,false,true,false
 2025-12-09,kubernetesUnifiedStorageQuotas,experimental,@grafana/search-and-storage,false,true,false
-2025-10-17,kubernetesLogsDrilldown,experimental,@grafana/observability-logs,false,true,false
+2025-10-16,kubernetesLogsDrilldown,experimental,@grafana/observability-logs,false,true,false
 2025-10-20,kubernetesQueryCaching,experimental,@grafana/grafana-operator-experience-squad,false,true,false
 2025-04-11,dashboardDisableSchemaValidationV1,experimental,@grafana/grafana-app-platform-squad,false,false,false
 2025-04-11,dashboardDisableSchemaValidationV2,experimental,@grafana/grafana-app-platform-squad,false,false,false
@@ -100,8 +100,8 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2025-07-31,useScopeSingleNodeEndpoint,experimental,@grafana/grafana-operator-experience-squad,false,false,true
 2025-08-29,useMultipleScopeNodesEndpoint,experimental,@grafana/grafana-operator-experience-squad,false,false,true
 2024-11-11,logQLScope,privatePreview,@grafana/oss-big-tent,false,false,false
-2024-02-28,sqlExpressions,preview,@grafana/grafana-datasources-core-services,false,false,false
-2025-07-24,sqlExpressionsColumnAutoComplete,experimental,@grafana/datapro,false,false,true
+2024-02-27,sqlExpressions,preview,@grafana/grafana-datasources-core-services,false,false,false
+2025-07-23,sqlExpressionsColumnAutoComplete,experimental,@grafana/datapro,false,false,true
 2024-02-12,kubernetesAggregator,experimental,@grafana/grafana-app-platform-squad,false,true,false
 2025-05-15,kubernetesAggregatorCapTokenAuth,experimental,@grafana/grafana-app-platform-squad,false,true,false
 2024-02-14,groupByVariable,experimental,@grafana/dashboards-squad,false,false,false
@@ -151,11 +151,11 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2024-10-17,unifiedStorageBigObjectsSupport,experimental,@grafana/search-and-storage,false,false,false
 2024-10-22,timeRangeProvider,experimental,@grafana/grafana-frontend-platform,false,false,false
 2025-11-05,timeRangePan,experimental,@grafana/dataviz-squad,false,false,true
-2025-11-21,newTimeRangeZoomShortcuts,experimental,@grafana/dataviz-squad,false,false,true
+2025-11-20,newTimeRangeZoomShortcuts,experimental,@grafana/dataviz-squad,false,false,true
 2024-10-24,azureMonitorDisableLogLimit,GA,@grafana/partner-datasources,false,false,false
 2024-12-20,playlistsReconciler,experimental,@grafana/grafana-app-platform-squad,false,true,false
 2024-11-14,passwordlessMagicLinkAuthentication,experimental,@grafana/identity-access-team,false,false,false
-2024-12-19,prometheusSpecialCharsInLabelValues,experimental,@grafana/oss-big-tent,false,false,true
+2024-12-18,prometheusSpecialCharsInLabelValues,experimental,@grafana/oss-big-tent,false,false,true
 2024-11-05,enableExtensionsAdminPage,experimental,@grafana/plugins-platform-backend,false,true,false
 2024-11-07,enableSCIM,preview,@grafana/identity-access-team,false,false,false
 2024-11-12,crashDetection,experimental,@grafana/observability-traces-and-profiling,false,false,true
@@ -170,7 +170,7 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2025-07-16,alertingAIAnalyzeCentralStateHistory,experimental,@grafana/alerting-squad,false,false,false
 2024-11-22,alertingNotificationsStepMode,GA,@grafana/alerting-squad,false,false,true
 2024-12-19,unifiedStorageSearchUI,experimental,@grafana/search-and-storage,false,false,false
-2024-12-13,elasticsearchCrossClusterSearch,GA,@grafana/partner-datasources,false,false,false
+2024-12-12,elasticsearchCrossClusterSearch,GA,@grafana/partner-datasources,false,false,false
 2024-12-13,lokiLabelNamesQueryApi,GA,@grafana/oss-big-tent,false,false,false
 2024-12-27,k8SFolderCounts,experimental,@grafana/search-and-storage,false,false,false
 2025-01-09,improvedExternalSessionHandlingSAML,GA,@grafana/identity-access-team,false,false,false
@@ -227,7 +227,7 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2025-08-01,alertEnrichmentConditional,experimental,@grafana/alerting-squad,false,false,false
 2025-06-10,alertingImportAlertmanagerAPI,experimental,@grafana/alerting-squad,false,false,false
 2025-08-01,alertingImportAlertmanagerUI,experimental,@grafana/alerting-squad,false,false,false
-2025-07-16,sharingDashboardImage,GA,@grafana/sharing-squad,false,false,true
+2025-07-15,sharingDashboardImage,GA,@grafana/sharing-squad,false,false,true
 2025-06-17,preferLibraryPanelTitle,privatePreview,@grafana/dashboards-squad,false,false,false
 2025-06-24,tabularNumbers,GA,@grafana/grafana-frontend-platform,false,false,false
 2025-06-25,newInfluxDSConfigPageDesign,privatePreview,@grafana/partner-datasources,false,false,false
@@ -256,7 +256,7 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2025-10-20,newGauge,preview,@grafana/dataviz-squad,false,false,true
 2025-11-12,newVizSuggestions,preview,@grafana/dataviz-squad,false,false,true
 2025-12-02,externalVizSuggestions,experimental,@grafana/dataviz-squad,false,false,true
-2025-12-19,heatmapRowsAxisOptions,experimental,@grafana/dataviz-squad,false,false,true
+2025-12-18,heatmapRowsAxisOptions,experimental,@grafana/dataviz-squad,false,false,true
 2025-10-17,preventPanelChromeOverflow,preview,@grafana/grafana-frontend-platform,false,false,true
 2025-10-31,jaegerEnableGrpcEndpoint,experimental,@grafana/oss-big-tent,false,false,false
 2025-10-17,pluginStoreServiceLoading,experimental,@grafana/plugins-platform-backend,false,false,false
@@ -279,3 +279,4 @@ Created,Name,Stage,Owner,requiresDevMode,RequiresRestart,FrontendOnly
 2026-01-06,secretsManagementAppPlatformAwsKeeper,experimental,@grafana/grafana-operator-experience-squad,false,false,false
 2026-01-07,profilesExemplars,experimental,@grafana/observability-traces-and-profiling,false,false,false
 2026-01-14,alertingSyncDispatchTimer,experimental,@grafana/alerting-squad,false,true,false
+2026-01-15,queryWithAssistant,experimental,@grafana/oss-big-tent,false,false,true

pkg/services/featuremgmt/toggles_gen.json

@@ -3708,6 +3708,24 @@
         "expression": "false"
       }
     },
+    {
+      "metadata": {
+        "name": "queryWithAssistant",
+        "resourceVersion": "1768832883692",
+        "creationTimestamp": "2026-01-15T12:40:17Z",
+        "deletionTimestamp": "2026-01-19T14:25:13Z",
+        "annotations": {
+          "grafana.app/updatedTimestamp": "2026-01-19 14:28:03.692934 +0000 UTC"
+        }
+      },
+      "spec": {
+        "description": "Enables the Query with Assistant button in the query editor",
+        "stage": "experimental",
+        "codeowner": "@grafana/oss-big-tent",
+        "frontend": true,
+        "expression": "false"
+      }
+    },
     {
       "metadata": {
         "name": "recentlyViewedDashboards",

public/app/features/query/components/QueryActionAssistantButton.test.tsx

@@ -0,0 +1,116 @@
+import { render, screen } from '@testing-library/react';
+
+import { AssistantHook, useAssistant } from '@grafana/assistant';
+import { CoreApp, DataSourceInstanceSettings } from '@grafana/data';
+import { DataQuery } from '@grafana/schema';
+
+import { QueryActionAssistantButton } from './QueryActionAssistantButton';
+// Mock the assistant hook
+jest.mock('@grafana/assistant', () => ({
+  useAssistant: jest.fn(),
+  createAssistantContextItem: jest.fn(),
+}));
+
+// Mock the runtime services that assistant depends on
+const mockConfig = {
+  featureToggles: {
+    queryWithAssistant: false,
+  },
+};
+
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  usePluginLinks: jest.fn().mockReturnValue({ links: [], isLoading: false }),
+  get config() {
+    return mockConfig;
+  },
+}));
+
+const useAssistantMock = jest.mocked(useAssistant);
+
+const mockDataSourceInstance: DataSourceInstanceSettings = {
+  uid: 'test-uid',
+  name: 'Test Datasource',
+  type: 'loki',
+} as DataSourceInstanceSettings;
+
+const mockQuery: DataQuery = {
+  refId: 'A',
+};
+
+const mockQueries: DataQuery[] = [mockQuery];
+
+const defaultProps = {
+  query: mockQuery,
+  queries: mockQueries,
+  dataSourceInstanceSettings: mockDataSourceInstance,
+  app: CoreApp.Explore,
+  datasourceApi: null,
+};
+
+describe('QueryActionAssistantButton', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    // Default: feature toggle enabled, assistant available
+    mockConfig.featureToggles.queryWithAssistant = true;
+    useAssistantMock.mockReturnValue({
+      isAvailable: true,
+      openAssistant: jest.fn(),
+    } as unknown as AssistantHook);
+  });
+
+  it('should render nothing when feature toggle is disabled', () => {
+    mockConfig.featureToggles.queryWithAssistant = false;
+    useAssistantMock.mockReturnValue({
+      isAvailable: true,
+      openAssistant: jest.fn(),
+    } as unknown as AssistantHook);
+
+    const { container } = render(<QueryActionAssistantButton {...defaultProps} />);
+
+    expect(container.firstChild).toBeNull();
+  });
+
+  it('should render nothing when app is not Explore, Dashboard, or PanelEditor', () => {
+    const { container } = render(<QueryActionAssistantButton {...defaultProps} app={CoreApp.Unknown} />);
+    expect(container.firstChild).toBeNull();
+  });
+
+  it('should render nothing when Assistant is not available', () => {
+    mockConfig.featureToggles.queryWithAssistant = true;
+    useAssistantMock.mockReturnValue({
+      isAvailable: false,
+      openAssistant: undefined,
+    } as unknown as AssistantHook);
+
+    const { container } = render(<QueryActionAssistantButton {...defaultProps} />);
+
+    expect(container.firstChild).toBeNull();
+  });
+
+  it('should render nothing when openAssistant is not provided', () => {
+    mockConfig.featureToggles.queryWithAssistant = true;
+    useAssistantMock.mockReturnValue({
+      isAvailable: true,
+      openAssistant: undefined,
+    } as unknown as AssistantHook);
+
+    const { container } = render(<QueryActionAssistantButton {...defaultProps} />);
+
+    expect(container.firstChild).toBeNull();
+  });
+
+  it('should render button when feature toggle is enabled and assistant is available', () => {
+    mockConfig.featureToggles.queryWithAssistant = true;
+    const mockOpenAssistant = jest.fn();
+    useAssistantMock.mockReturnValue({
+      isAvailable: true,
+      openAssistant: mockOpenAssistant,
+    } as unknown as AssistantHook);
+
+    render(<QueryActionAssistantButton {...defaultProps} />);
+
+    const button = screen.getByRole('button', { name: /query with assistant/i });
+    expect(button).toBeInTheDocument();
+  });
+});

public/app/features/query/components/QueryActionAssistantButton.tsx

@@ -0,0 +1,154 @@
+import { useAssistant, createAssistantContextItem } from '@grafana/assistant';
+import { CoreApp, DataSourceApi, DataSourceInstanceSettings } from '@grafana/data';
+import { t } from '@grafana/i18n';
+import { config } from '@grafana/runtime';
+import { DataQuery, DataSourceJsonData } from '@grafana/schema';
+import { Button } from '@grafana/ui';
+import { queryIsEmpty } from 'app/core/utils/query';
+
+interface QueryActionAssistantButtonProps<TQuery extends DataQuery = DataQuery> {
+  query: TQuery;
+  queries: TQuery[];
+  dataSourceInstanceSettings: DataSourceInstanceSettings;
+  app?: CoreApp;
+  datasourceApi: DataSourceApi<TQuery, DataSourceJsonData, {}> | null;
+}
+
+export function QueryActionAssistantButton<TQuery extends DataQuery = DataQuery>({
+  query,
+  queries,
+  dataSourceInstanceSettings,
+  app,
+  datasourceApi,
+}: QueryActionAssistantButtonProps<TQuery>) {
+  const { isAvailable, openAssistant } = useAssistant();
+
+  // Check if the feature toggle is enabled
+  if (!config.featureToggles.queryWithAssistant) {
+    return null;
+  }
+
+  if (!isAvailable || !openAssistant) {
+    return null;
+  }
+
+  // Only show for Explore and Dashboard apps
+  if (app !== CoreApp.Explore && app !== CoreApp.Dashboard && app !== CoreApp.PanelEditor) {
+    return null;
+  }
+
+  // Only show for loki and prometheus datasources
+  const pluginId = dataSourceInstanceSettings.type;
+  if (pluginId !== 'loki' && pluginId !== 'prometheus') {
+    return null;
+  }
+  const origin = `grafana/query-editor/${pluginId}/${app ?? CoreApp.Unknown}`;
+
+  // Check if current query has content
+  const hasCurrentQuery = !queryIsEmpty(query);
+  const otherQueries = queries.filter((q) => q.refId !== query.refId && !queryIsEmpty(q));
+
+  // Build context items
+  const context = [
+    createAssistantContextItem('datasource', {
+      datasourceUid: dataSourceInstanceSettings.uid,
+    }),
+  ];
+
+  // Add current query if it has content
+  if (hasCurrentQuery) {
+    context.push(
+      createAssistantContextItem('structured', {
+        title: t('query-operation.header.current-query', 'Current query'),
+        data: query,
+      })
+    );
+  }
+
+  // Add other queries if they exist
+  if (otherQueries.length > 0) {
+    context.push(
+      createAssistantContextItem('structured', {
+        title: t('query-operation.header.other-queries', 'Other queries'),
+        data: {
+          queries: otherQueries,
+        },
+      })
+    );
+  }
+
+  // Get query display text to determine if we're creating or updating
+  const queryDisplayText =
+    hasCurrentQuery && datasourceApi?.getQueryDisplayText ? datasourceApi.getQueryDisplayText(query) : null;
+
+  // Determine if we're creating or updating based on queryDisplayText
+  const isUpdating = !!queryDisplayText;
+  const actionText = isUpdating
+    ? t(
+        'query-operation.header.assistant-prompt-update',
+        'Help me update the current query to answer my questions and provide the insights I need.'
+      )
+    : t(
+        'query-operation.header.assistant-prompt-create',
+        'Help me create a new query to answer my questions and provide the insights I need.'
+      );
+
+  // Format app name nicely
+  const appName =
+    app === CoreApp.Explore
+      ? t('query-operation.header.app-explore', 'Explore')
+      : app === CoreApp.Dashboard
+        ? t('query-operation.header.app-dashboard', 'Dashboard')
+        : '';
+
+  // Build the prompt with proper formatting
+  const codeBlockLines: string[] = [];
+
+  if (queryDisplayText) {
+    codeBlockLines.push(t('query-operation.header.current-query-label', 'Current query:') + ` ${queryDisplayText}`);
+  }
+
+  codeBlockLines.push(
+    t('query-operation.header.selected-datasource-label', 'Selected data source:') +
+      ` ${dataSourceInstanceSettings.name}`
+  );
+
+  if (appName) {
+    codeBlockLines.push(t('query-operation.header.app-label', 'App:') + ` ${appName}`);
+  }
+
+  // Add actionable sentence to motivate users
+  const actionableSentence = isUpdating
+    ? t(
+        'query-operation.header.assistant-actionable-update',
+        'Please describe what you want to change or improve in this query.'
+      )
+    : t(
+        'query-operation.header.assistant-actionable-create',
+        "Please describe what you want to query and what insights you're looking for."
+      );
+
+  // Build final prompt with code block
+  const prompt = [actionText, '```', ...codeBlockLines, '```', actionableSentence].join('\n');
+
+  const handleClick = () => {
+    openAssistant({
+      origin,
+      prompt,
+      context,
+      autoSend: false,
+    });
+  };
+
+  return (
+    <Button
+      size="sm"
+      variant="secondary"
+      icon="ai"
+      onClick={handleClick}
+      title={t('query-operation.header.query-with-assistant', 'Query with Assistant')}
+    >
+      {t('query-operation.header.query-with-assistant', 'Query with Assistant')}
+    </Button>
+  );
+}

public/app/features/query/components/QueryEditorRow.tsx

@@ -36,6 +36,7 @@ import {
 import { useQueryLibraryContext } from '../../explore/QueryLibrary/QueryLibraryContext';
 import { ExpressionDatasourceUID } from '../../expressions/types';
 
+import { QueryActionAssistantButton } from './QueryActionAssistantButton';
 import { QueryActionComponent, RowActionComponents } from './QueryActionComponent';
 import { QueryEditorRowHeader } from './QueryEditorRowHeader';
 import { QueryErrorAlert } from './QueryErrorAlert';
@@ -467,6 +468,7 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
 
   renderHeader = (props: QueryOperationRowRenderProps) => {
     const { app, query, dataSource, onChangeDataSource, onChange, queries, renderHeaderExtras, hideRefId } = this.props;
+    const { datasource } = this.state;
 
     return (
       <QueryEditorRowHeader
@@ -477,7 +479,18 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
         hidden={query.hide}
         onChange={onChange}
         collapsedText={!props.isOpen ? this.renderCollapsedText() : null}
-        renderExtras={renderHeaderExtras}
+        renderExtras={() => (
+          <>
+            <QueryActionAssistantButton
+              query={query}
+              queries={queries}
+              dataSourceInstanceSettings={this.props.dataSource}
+              datasourceApi={datasource}
+              app={app}
+            />
+            {renderHeaderExtras && renderHeaderExtras()}
+          </>
+        )}
         alerting={app === CoreApp.UnifiedAlerting}
         hideRefId={hideRefId}
       />

public/locales/en-US/grafana.json

@@ -12544,13 +12544,25 @@
   },
   "query-operation": {
     "header": {
+      "app-dashboard": "Dashboard",
+      "app-explore": "Explore",
+      "app-label": "App:",
+      "assistant-actionable-create": "Please describe what you want to query and what insights you're looking for.",
+      "assistant-actionable-update": "Please describe what you want to change or improve in this query.",
+      "assistant-prompt-create": "Help me create a new query to answer my questions and provide the insights I need.",
+      "assistant-prompt-update": "Help me update the current query to answer my questions and provide the insights I need.",
       "collapse-row": "Collapse query row",
+      "current-query": "Current query",
+      "current-query-label": "Current query:",
       "datasource-help": "Show data source help",
       "drag-and-drop": "Drag and drop to reorder",
       "duplicate-query": "Duplicate query",
       "expand-row": "Expand query row",
       "hide-response": "Hide response",
+      "other-queries": "Other queries",
+      "query-with-assistant": "Query with Assistant",
       "remove-query": "Remove query",
+      "selected-datasource-label": "Selected data source:",
       "show-response": "Show response"
     },
     "query-editor-not-exported": "Data source plugin does not export any Query Editor component"

public/test/mocks/assistant.ts

@@ -0,0 +1,19 @@
+// Mock for @grafana/assistant to prevent initialization errors in tests
+// The real module tries to call getObservablePluginLinks() during initialization
+// which fails because Grafana hasn't started. This mock prevents that.
+
+export const useAssistant = jest.fn().mockReturnValue({
+  isAvailable: false,
+  openAssistant: undefined,
+  closeAssistant: jest.fn(),
+  toggleAssistant: jest.fn(),
+});
+
+export const createAssistantContextItem = jest.fn();
+
+// Additional exports that may be used
+export const toggleAssistant = jest.fn();
+export const isAssistantAvailable = jest.fn().mockReturnValue(false);
+
+// Type exports (if needed
+export type AssistantHook = ReturnType<typeof useAssistant>;