feat(explorer): let the assistant run a query on the metric explorer

src/pages/explorer/Prometheus/index.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useState } from 'react';
+import React, { useEffect, useRef, useState } from 'react';
 import { useLocation, useHistory } from 'react-router-dom';
 import queryString from 'query-string';
 import moment from 'moment';
@@ -14,10 +14,12 @@ import { getHistoryEventsById } from '@/services/warning';
 
 import { AiButton } from '@/components/AiChatNG/FlashAiButton';
 import { buildPageFrom, getExplorerPrompts } from '@/components/AiChatNG/recommend';
+import { useAiChatContext } from '@/components/AiChatNG';
 
 import { queryStringOptions } from '../constants';
 import ProbeBanner from '../components/ProbeBanner';
 import HistoricalRecords, { setLocalQueryHistory } from './HistoricalRecords';
+import { useMetricExplorerAIActions } from './useMetricExplorerAIActions';
 
 const LOCAL_KEY = 'n9e-query-promql-history';
 
@@ -73,6 +75,24 @@ export default function Prometheus(props: IProps) {
   const [promql, setPromql] = useState<string>(defaultPromQL);
   // 体检落地横幅：仅 __from=ds_verify 进入且首个面板展示；用户接管（改查询/点查询）或点 × 后收起
   const [probeBannerVisible, setProbeBannerVisible] = useState<boolean>(query.__from === 'ds_verify' && panelIdx === 0);
+  const { queryPageFrom } = useAiChatContext();
+  // Scopes the lookup below to this panel. Panels on this page can each be on a
+  // different data source, and only some of them render a PromQL box at all, so
+  // there is no position in a page-wide query that reliably means "this one".
+  // `display: contents` keeps the wrapper out of the layout entirely.
+  const panelRef = useRef<HTMLDivElement>(null);
+
+  useMetricExplorerAIActions({
+    // The chat is a single panel at the app root, so at most one explorer panel
+    // is ever the one being talked about. Whichever panel's AI button opened it
+    // put its own key on page_from; that is the one whose query box may be
+    // written to.
+    enabled: Boolean(panelKey) && (queryPageFrom?.param?.panelKey as string | undefined) === panelKey,
+    datasourceValue,
+    setPromql,
+    setTimeRange: setDefaultTimeState,
+    getQueryInput: () => panelRef.current?.querySelector('.prom-graph-expression-input-ng') ?? null,
+  });
 
   useEffect(() => {
     if (query.__event_id) {
@@ -104,7 +124,7 @@ export default function Prometheus(props: IProps) {
   }, []);
 
   return (
-    <>
+    <div ref={panelRef} style={{ display: 'contents' }}>
       <PromGraph
         // key={promql} // 当存在 query.__event_id 时需要异步获取 datasourceValue 和 prom_ql，这时需要强制重新渲染
         type={query.mode as IMode}
@@ -175,6 +195,10 @@ export default function Prometheus(props: IProps) {
                 param: {
                   datasource_type: 'prometheus',
                   datasource_id: datasourceValue,
+                  // Says which panel the conversation belongs to. Metric.tsx
+                  // closes the chat when that panel is removed, and the panel
+                  // registers its page action only while it is the owner.
+                  panelKey,
                 },
               })}
               queryAction={{
@@ -194,6 +218,6 @@ export default function Prometheus(props: IProps) {
         }
         showExportButton
       />
-    </>
+    </div>
   );
 }

src/pages/explorer/Prometheus/useMetricExplorerAIActions.test.ts

@@ -0,0 +1,163 @@
+/** @jest-environment jsdom */
+import { renderHook } from '@testing-library/react';
+import type { ActionRunContext, UIAction } from '@flashcatcloud/ai-kit/actions';
+
+import { useMetricExplorerAIActions, MetricExplorerAIActionsOptions } from './useMetricExplorerAIActions';
+
+/**
+ * The registry is faked rather than imported, because ai-kit ships ESM that
+ * this project's jest transform does not cover. Its schema validation and its
+ * unsupported/timeout handling are the package's own contract and are tested
+ * there; what is worth pinning here is only what this hook decides — when it
+ * registers at all, and what it does to the page when the action runs.
+ */
+const registered = new Map<string, UIAction<any>>();
+const register = jest.fn((actions: UIAction<any>[]) => {
+  actions.forEach((action) => registered.set(action.name, action));
+  return () => actions.forEach((action) => registered.delete(action.name));
+});
+
+jest.mock('@/components/AiChatNG/uiActionRuntime', () => ({
+  uiActionRuntime: {
+    register: (actions: UIAction<any>[], page: unknown) => register(actions, page),
+  },
+}));
+
+const ACTION = 'set_metric_query';
+
+function options(overrides: Partial<MetricExplorerAIActionsOptions> = {}): MetricExplorerAIActionsOptions {
+  return {
+    enabled: true,
+    datasourceValue: 18001,
+    setPromql: jest.fn(),
+    setTimeRange: jest.fn(),
+    getQueryInput: () => null,
+    ...overrides,
+  };
+}
+
+function runContext(): ActionRunContext {
+  return {
+    callId: 'call-1',
+    signal: new AbortController().signal,
+    feedback: {
+      reveal: jest.fn().mockResolvedValue(undefined),
+      highlight: jest.fn(),
+      moveCursor: jest.fn().mockResolvedValue(undefined),
+      click: jest.fn().mockResolvedValue(undefined),
+      setControlledRegion: jest.fn(),
+      clear: jest.fn(),
+    },
+  };
+}
+
+function run(args: Record<string, unknown>) {
+  const action = registered.get(ACTION);
+  if (!action) throw new Error(`${ACTION} is not registered`);
+  return action.run(args as never, runContext());
+}
+
+beforeEach(() => {
+  registered.clear();
+  register.mockClear();
+});
+
+describe('useMetricExplorerAIActions', () => {
+  it('registers while the panel owns the conversation and drops it on unmount', () => {
+    const { unmount } = renderHook(() => useMetricExplorerAIActions(options()));
+    expect(registered.has(ACTION)).toBe(true);
+    unmount();
+    expect(registered.has(ACTION)).toBe(false);
+  });
+
+  it('registers nothing for a panel that does not own the conversation', () => {
+    renderHook(() => useMetricExplorerAIActions(options({ enabled: false })));
+    expect(register).not.toHaveBeenCalled();
+    expect(registered.has(ACTION)).toBe(false);
+  });
+
+  it('follows the conversation when it moves to this panel', () => {
+    const { rerender, unmount } = renderHook((props: MetricExplorerAIActionsOptions) => useMetricExplorerAIActions(props), {
+      initialProps: options({ enabled: false }),
+    });
+    expect(registered.has(ACTION)).toBe(false);
+    rerender(options({ enabled: true }));
+    expect(registered.has(ACTION)).toBe(true);
+    unmount();
+  });
+
+  it('does not re-register when only the query or the data source changes', () => {
+    const { rerender, unmount } = renderHook((props: MetricExplorerAIActionsOptions) => useMetricExplorerAIActions(props), {
+      initialProps: options(),
+    });
+    expect(register).toHaveBeenCalledTimes(1);
+    rerender(options({ datasourceValue: 999 }));
+    expect(register).toHaveBeenCalledTimes(1);
+    unmount();
+  });
+
+  it('writes the expression into the panel and leaves the time range alone', async () => {
+    const setPromql = jest.fn();
+    const setTimeRange = jest.fn();
+    renderHook(() => useMetricExplorerAIActions(options({ setPromql, setTimeRange })));
+
+    const result = (await run({ promql: '  cpu_usage_active  ' })) as { promql: string };
+
+    // Trimmed: a leading space in the box is something the user has to clean up.
+    expect(setPromql).toHaveBeenCalledWith('cpu_usage_active');
+    expect(result.promql).toBe('cpu_usage_active');
+    // Not asked for, so the panel keeps the window the user was reading.
+    expect(setTimeRange).not.toHaveBeenCalled();
+  });
+
+  it('reads the live data source rather than the one captured at registration', async () => {
+    const { rerender } = renderHook((props: MetricExplorerAIActionsOptions) => useMetricExplorerAIActions(props), {
+      initialProps: options({ datasourceValue: 1 }),
+    });
+    rerender(options({ datasourceValue: 2 }));
+
+    const result = (await run({ promql: 'up' })) as { datasource_id: number };
+
+    expect(result.datasource_id).toBe(2);
+  });
+
+  it('moves the time range when one is asked for', async () => {
+    const setTimeRange = jest.fn();
+    renderHook(() => useMetricExplorerAIActions(options({ setTimeRange })));
+
+    await run({ promql: 'up', time_range: { start: 'now-6h', end: 'now' } });
+
+    expect(setTimeRange).toHaveBeenCalledWith({ start: 'now-6h', end: 'now' });
+  });
+
+  it('treats an explicit null window as no window', async () => {
+    const setTimeRange = jest.fn();
+    renderHook(() => useMetricExplorerAIActions(options({ setTimeRange })));
+
+    // The runtime's validator lets null through where it would reject a
+    // half-filled object, so this is the one malformed window run() still sees.
+    await run({ promql: 'up', time_range: null });
+
+    expect(setTimeRange).not.toHaveBeenCalled();
+  });
+
+  it('points the cursor at the panel that owns the action, not a page-wide guess', async () => {
+    const input = document.createElement('div');
+    const context = runContext();
+    renderHook(() => useMetricExplorerAIActions(options({ getQueryInput: () => input })));
+
+    const action = registered.get(ACTION)!;
+    await action.run({ promql: 'up' } as never, context);
+
+    expect(context.feedback.moveCursor).toHaveBeenCalledWith(input);
+    expect(context.feedback.highlight).toHaveBeenCalledWith(input);
+  });
+
+  it('refuses a blank expression instead of clearing the box', async () => {
+    const setPromql = jest.fn();
+    renderHook(() => useMetricExplorerAIActions(options({ setPromql })));
+
+    await expect(run({ promql: '   ' })).rejects.toThrow();
+    expect(setPromql).not.toHaveBeenCalled();
+  });
+});

src/pages/explorer/Prometheus/useMetricExplorerAIActions.ts

@@ -0,0 +1,134 @@
+import { useEffect, useRef } from 'react';
+
+import { uiActionRuntime } from '@/components/AiChatNG/uiActionRuntime';
+import type { UIAction } from '@flashcatcloud/ai-kit/actions';
+import type { IRawTimeRange } from '@/components/TimeRangePicker';
+
+/**
+ * The ad-hoc metric query page, described to the assistant.
+ *
+ * The assistant could already write PromQL into a reply, but the user then had
+ * to select it, copy it and paste it into the box — and a query the assistant
+ * cannot run is a query it never sees the result of. This action closes that
+ * gap: the expression lands in the same input the user types into, and the
+ * page's own effects re-run the query from there, exactly as they do after a
+ * keystroke.
+ *
+ * It sets the query and nothing else. There is no "save" on this page, so
+ * there is nothing here to undo — the previous expression is one Ctrl+Z or one
+ * click of 历史记录 away.
+ */
+
+export interface MetricExplorerAIActionsOptions {
+  /** Whether this panel owns the open conversation. See the note in the hook. */
+  enabled: boolean;
+  /** The data source the panel is pointed at, echoed back for the card. */
+  datasourceValue: number;
+  /** Writes the expression into the panel, which re-runs the query. */
+  setPromql: (promql: string) => void;
+  /** Moves the panel's time range. Only used when the model asks for one. */
+  setTimeRange: (range: IRawTimeRange) => void;
+  /**
+   * This panel's own PromQL box, used only to point the cursor at what changed.
+   *
+   * Resolved by the caller rather than looked up here: panels on this page can
+   * each be on a different data source and only some of them render a PromQL
+   * box, so no page-wide position or selector reliably means "this panel".
+   */
+  getQueryInput: () => Element | null;
+}
+
+interface SetMetricQueryArgs {
+  promql: string;
+  /** Both ends, or nothing: the schema's nested `required` is enforced. Null
+   *  still gets through, which is the one case `run()` has to handle. */
+  time_range?: { start: string; end: string } | null;
+}
+
+export function useMetricExplorerAIActions(options: MetricExplorerAIActionsOptions): void {
+  const latest = useRef(options);
+  latest.current = options;
+
+  useEffect(() => {
+    // Every panel on this page is the same component, so without this guard a
+    // second panel would register the same action name and quietly take over
+    // the first one's — ai-kit lets the later registration win. Registering
+    // only from the panel whose AI button opened the conversation keeps
+    // "which panel does this write to" answerable rather than incidental.
+    if (!options.enabled) return;
+
+    const actions: UIAction<SetMetricQueryArgs>[] = [
+      {
+        name: 'set_metric_query',
+        description:
+          'Put a PromQL expression into the ad-hoc metric query page and run it. ' +
+          'Use this instead of only printing the expression in the reply: it fills the same input the user types into, so the chart and the series table refresh immediately. ' +
+          'Send the expression you have already verified against this data source — this action runs it, it does not check it. ' +
+          'It changes nothing else: no dashboard, no alert rule, nothing is saved.',
+        schema: {
+          type: 'object',
+          properties: {
+            promql: {
+              type: 'string',
+              description: 'The complete expression to run, exactly as it should appear in the input box.',
+            },
+            time_range: {
+              type: 'object',
+              description:
+                'Optional. Only send it when the user asked for a particular window, otherwise the panel keeps the range it is on. ' +
+                'Both ends accept the relative syntax the page itself uses, e.g. "now-6h" and "now".',
+              properties: {
+                start: { type: 'string', description: 'Window start, e.g. "now-6h".' },
+                end: { type: 'string', description: 'Window end, usually "now".' },
+              },
+              required: ['start', 'end'],
+            },
+          },
+          required: ['promql'],
+        },
+        // Nothing is persisted, and the user already clicked the card to get
+        // here. A second confirmation would be a dialog in front of a dialog.
+        policy: 'auto',
+        run: async (args, ctx) => {
+          const promql = args.promql?.trim();
+          // The runtime already rejects a missing or empty string. Whitespace
+          // that only looks like a query is the part it cannot catch.
+          if (!promql) throw new Error('No expression was supplied.');
+
+          const { datasourceValue, setPromql, setTimeRange, getQueryInput } = latest.current;
+          const anchor = getQueryInput();
+          await ctx.feedback.moveCursor(anchor);
+          ctx.feedback.highlight(anchor);
+
+          // Half a window would silently reframe the chart around a range nobody
+          // asked for, but the schema already rules that out — only an explicit
+          // null still reaches here, and it means "no window", same as absent.
+          const movedRange = args.time_range ?? undefined;
+          if (movedRange) {
+            setTimeRange(movedRange);
+          }
+          setPromql(promql);
+
+          return {
+            promql,
+            datasource_id: datasourceValue,
+            time_range: movedRange,
+            // Said plainly because the card shows this to the user, who can see
+            // the page for themselves and would notice a grander claim.
+            note: 'Written into the query box and run. Nothing was saved.',
+          };
+        },
+      },
+    ];
+
+    return uiActionRuntime.register(actions, {
+      route: `${window.location.pathname}${window.location.search}`,
+      title: '指标分析 · 即时查询',
+      summary:
+        'The user is on the ad-hoc metric query page, running PromQL against one data source and reading the result as a chart or a series table. ' +
+        'Nothing on this page is saved, so a query can be replaced freely.',
+    });
+    // Live values are read through a ref so that typing in the box, or moving
+    // the time range, does not churn the registration.
+  }, [options.enabled]);
+}