[BUGFIX] Sync dashboard state with query parameters

dashboards/src/context/DashboardProvider/DashboardProvider.tsx

@@ -95,6 +95,7 @@ export function DashboardProvider(props: DashboardProviderProps): ReactElement {
   const { defaultPluginKinds } = usePluginRegistry();
   const defaultPanelKind = defaultPluginKinds?.['Panel'] ?? '';
   const { data: plugin } = usePlugin('Panel', defaultPanelKind);
+  const { viewPanelRef } = props.initialState;
 
   useEffect(() => {
     if (plugin === undefined) return;
@@ -107,6 +108,10 @@ export function DashboardProvider(props: DashboardProviderProps): ReactElement {
     });
   }, [plugin, store, defaultPanelKind]);
 
+  useEffect(() => {
+    store.getState().setViewPanelFromRef(viewPanelRef);
+  }, [store, viewPanelRef]);
+
   return (
     <DashboardContext.Provider value={store as StoreApi<DashboardStoreState>}>
       {props.children}

dashboards/src/context/DashboardProvider/view-panel-slice.ts

@@ -32,6 +32,7 @@ export interface ViewPanelSlice {
   viewPanel: ViewPanelState;
   getViewPanel: () => PanelGroupItemId | undefined;
   setViewPanel: (panelGroupItemId?: PanelGroupItemId) => void;
+  setViewPanelFromRef: (panelRef?: VirtualPanelRef) => void;
 }
 
 export interface ViewPanelState {
@@ -45,7 +46,7 @@ export interface ViewPanelState {
  */
 export function createViewPanelSlice(
   viewPanelRef?: VirtualPanelRef,
-  setViewPanelRef?: (ref: VirtualPanelRef | undefined) => void
+  setViewPanelRefQueryParam?: (ref: VirtualPanelRef | undefined) => void
 ): StateCreator<ViewPanelSlice & PanelGroupSlice, Middleware, [], ViewPanelSlice> {
   return (set, get) => ({
     viewPanel: {
@@ -64,14 +65,37 @@ export function createViewPanelSlice(
           panelGroupItemId: panelGroupItemId,
         };
         const panelRef = findPanelRefOfPanelGroupItemId(get().panelGroups, panelGroupItemId);
-        if (setViewPanelRef) {
-          setViewPanelRef(panelRef);
+        if (setViewPanelRefQueryParam) {
+          setViewPanelRefQueryParam(panelRef);
         }
       });
     },
+
+    setViewPanelFromRef(panelRef?: VirtualPanelRef): void {
+      set((state) => {
+        const currentPanelRef =
+          state.viewPanel.panelRef ??
+          findPanelRefOfPanelGroupItemId(state.panelGroups, state.viewPanel.panelGroupItemId);
+        if (areViewPanelRefsEqual(currentPanelRef, panelRef)) {
+          return;
+        }
+        state.viewPanel = {
+          panelGroupItemId: undefined,
+          panelRef,
+        };
+      });
+    },
   });
 }
 
+function areViewPanelRefsEqual(left?: VirtualPanelRef, right?: VirtualPanelRef): boolean {
+  return (
+    left?.ref === right?.ref &&
+    left?.repeatVariable?.[0] === right?.repeatVariable?.[0] &&
+    left?.repeatVariable?.[1] === right?.repeatVariable?.[1]
+  );
+}
+
 function getViewPanelGroupId(
   panelGroups: Record<PanelGroupId, PanelGroupDefinition>,
   panelGroupItemId?: PanelGroupItemId,

dashboards/src/context/QueryParamSynchronization.test.tsx

@@ -0,0 +1,113 @@
+// Copyright The Perses Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { act, screen, waitFor } from '@testing-library/react';
+import { createMemoryHistory } from 'history';
+import { ReactElement } from 'react';
+import { TimeRangeProviderBasic } from '@perses-dev/plugin-system';
+import { VariableDefinition } from '@perses-dev/spec';
+import { createDashboardProviderSpy, getTestDashboard, renderWithContext } from '../test';
+import { DashboardProviderWithQueryParams } from './DashboardProvider/DashboardProviderWithQueryParams';
+import { useVariableDefinitionStates, VariableProviderWithQueryParams } from './VariableProvider';
+
+const variableDefinitions: VariableDefinition[] = [
+  {
+    kind: 'TextVariable',
+    spec: {
+      name: 'traceId',
+      value: 'default-trace',
+    },
+  },
+];
+
+function VariableValue(): ReactElement {
+  const variables = useVariableDefinitionStates(['traceId']);
+  return <div>{variables.traceId?.value}</div>;
+}
+
+describe('query parameter synchronization', () => {
+  test('updates variable values when navigation changes query parameters', async () => {
+    const history = createMemoryHistory({ initialEntries: ['/?var-traceId=first-trace'] });
+    renderWithContext(
+      <TimeRangeProviderBasic initialTimeRange={{ pastDuration: '30m' }}>
+        <VariableProviderWithQueryParams initialVariableDefinitions={variableDefinitions}>
+          <VariableValue />
+        </VariableProviderWithQueryParams>
+      </TimeRangeProviderBasic>,
+      undefined,
+      history
+    );
+
+    expect(await screen.findByText('first-trace')).toBeInTheDocument();
+
+    act(() => history.push('/?var-traceId=second-trace'));
+    expect(await screen.findByText('second-trace')).toBeInTheDocument();
+
+    act(() => history.push('/'));
+    expect(await screen.findByText('default-trace')).toBeInTheDocument();
+  });
+
+  test('updates the viewed panel when navigation changes query parameters', async () => {
+    const firstPanelRef = { ref: 'cpu' };
+    const secondPanelRef = { ref: 'memory', repeatVariable: ['instance', 'demo'] as [string, string] };
+    const history = createMemoryHistory({
+      initialEntries: [`/?viewPanelRef=${encodeURIComponent(JSON.stringify(firstPanelRef))}`],
+    });
+    const { DashboardProviderSpy, store } = createDashboardProviderSpy();
+
+    renderWithContext(
+      <DashboardProviderWithQueryParams initialState={{ dashboardResource: getTestDashboard() }}>
+        <DashboardProviderSpy />
+      </DashboardProviderWithQueryParams>,
+      undefined,
+      history
+    );
+
+    await waitFor(() => expect(store.value?.getState().viewPanel.panelRef).toEqual(firstPanelRef));
+
+    act(() => history.push(`/?viewPanelRef=${encodeURIComponent(JSON.stringify(secondPanelRef))}`));
+    await waitFor(() => expect(store.value?.getState().viewPanel.panelRef).toEqual(secondPanelRef));
+
+    act(() => history.push('/'));
+    await waitFor(() => expect(store.value?.getState().viewPanel.panelRef).toBeUndefined());
+  });
+
+  test('preserves the viewed panel when its store update changes query parameters', async () => {
+    const history = createMemoryHistory();
+    const { DashboardProviderSpy, store } = createDashboardProviderSpy();
+
+    renderWithContext(
+      <DashboardProviderWithQueryParams initialState={{ dashboardResource: getTestDashboard() }}>
+        <DashboardProviderSpy />
+      </DashboardProviderWithQueryParams>,
+      undefined,
+      history
+    );
+
+    const dashboardState = store.value?.getState();
+    const panelGroup = Object.values(dashboardState?.panelGroups ?? {})[0];
+    const panelGroupItemLayoutId = Object.keys(panelGroup?.itemPanelKeys ?? {})[0];
+    if (!panelGroup || !panelGroupItemLayoutId) {
+      throw new Error('Expected the test dashboard to contain a panel');
+    }
+
+    const panelGroupItemId = {
+      panelGroupId: panelGroup.id,
+      panelGroupItemLayoutId,
+    };
+    act(() => store.value?.getState().setViewPanel(panelGroupItemId));
+
+    await waitFor(() => expect(history.location.search).toContain('viewPanelRef='));
+    expect(store.value?.getState().viewPanel.panelGroupItemId).toEqual(panelGroupItemId);
+  });
+});

dashboards/src/context/VariableProvider/VariableProvider.tsx

@@ -11,7 +11,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { createContext, ReactElement, ReactNode, useContext, useMemo, useState } from 'react';
+import { createContext, ReactElement, ReactNode, useContext, useEffect, useMemo, useState } from 'react';
 import { createStore, StoreApi, useStore } from 'zustand';
 import { useStoreWithEqualityFn } from 'zustand/traditional';
 import { immer } from 'zustand/middleware/immer';
@@ -101,6 +101,7 @@ type VariableDefinitionStore = {
    */
   setVariableLoading: (name: VariableName, loading: boolean, source?: string) => void;
   setVariableDefinitions: (definitions: VariableDefinition[]) => void;
+  setVariableValuesFromQueryParams: (values: Record<string, VariableValue>) => void;
   setVariableDefaultValues: () => VariableDefinition[];
   getSavedVariablesStatus: () => { isSavedVariableModified: boolean; modifiedVariableNames: string[] };
 };
@@ -367,6 +368,36 @@ function createVariableDefinitionStore({
             '[Variables] setVariableDefinitions' // Used for action name in Redux devtools
           );
         },
+        setVariableValuesFromQueryParams(values: Record<string, VariableValue>): void {
+          set(
+            (state) => {
+              const hydratedState = hydrateVariableDefinitionStates(
+                state.variableDefinitions,
+                values,
+                state.externalVariableDefinitions
+              );
+
+              const updateValue = (name: string, source?: string): void => {
+                const currentState = state.variableState.get({ name, source });
+                const hydratedVariableState = hydratedState.get({ name, source });
+                if (
+                  currentState &&
+                  hydratedVariableState &&
+                  !areVariableValuesEqual(currentState.value, hydratedVariableState.value)
+                ) {
+                  currentState.value = hydratedVariableState.value;
+                }
+              };
+
+              state.variableDefinitions.forEach((definition) => updateValue(definition.spec.name));
+              state.externalVariableDefinitions.forEach(({ source, definitions }) => {
+                definitions.forEach((definition) => updateValue(definition.spec.name, source));
+              });
+            },
+            false,
+            '[Variables] setVariableValuesFromQueryParams'
+          );
+        },
         setVariableOptions(name, options, source?: string): void {
           set(
             (state) => {
@@ -518,10 +549,22 @@ export function VariableProviderWithQueryParams({
   const [store] = useState(() =>
     createVariableDefinitionStore({ initialVariableDefinitions, externalVariableDefinitions, queryParams })
   );
+  const queryParamValues = queryParams[0];
+
+  useEffect(() => {
+    store.getState().setVariableValuesFromQueryParams(getInitalValuesFromQueryParameters(queryParamValues));
+  }, [queryParamValues, store]);
 
   return (
     <VariableDefinitionStoreContext.Provider value={store}>
       <PluginProvider builtinVariables={builtinVariables}>{children}</PluginProvider>
     </VariableDefinitionStoreContext.Provider>
   );
 }
+
+function areVariableValuesEqual(left: VariableValue, right: VariableValue): boolean {
+  if (Array.isArray(left) && Array.isArray(right)) {
+    return left.length === right.length && left.every((value, index) => value === right[index]);
+  }
+  return left === right;
+}

dashboards/src/context/VariableProvider/hydrationUtils.ts

@@ -85,7 +85,7 @@ export function hydrateVariableDefinitionStates(
       externalDef.definitions.forEach((v) => {
         const name = v.spec.name;
         const param = initialValues[name];
-        const initialValue = param ? param : null;
+        const initialValue = param ?? null;
         state.set(
           { source, name },
           {
@@ -106,7 +106,7 @@ export function hydrateVariableDefinitionStates(
   localDefinitions.forEach((v) => {
     const name = v.spec.name;
     const param = initialValues[name];
-    const initialValue = param ? param : null;
+    const initialValue = param ?? null;
     state.set(
       { name },
       {