Initial profile state PoC implementation

Author: davismcphee

src/platform/plugins/shared/discover/public/application/main/state_management/redux/constants.ts

@@ -43,6 +43,7 @@ export const DEFAULT_TAB_STATE: Omit<TabState, keyof TabItem> = {
     fieldsToReset: 'none',
     snapshotsByProfileId: {},
   },
+  profileState: {},
   expandedDoc: undefined,
   expandedDocOwner: undefined,
   renderDocumentViewMeta: undefined,

src/platform/plugins/shared/discover/public/application/main/state_management/redux/context_awareness_toolkit.ts

@@ -7,17 +7,28 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import { distinctUntilChanged, from, map } from 'rxjs';
 import type { ContextAwarenessToolkit } from '../../../../context_awareness/toolkit';
 import type { InternalStateStore } from './internal_state';
 import { internalStateActions } from '.';
+import { selectTab } from './selectors';
+import type {
+  ProfileStateAdapter,
+  ProfileStateKey,
+  ProfileStateRegistry,
+} from '../../../../context_awareness';
 
 export const createContextAwarenessToolkit = ({
   internalState,
+  stateRegistry,
   tabId,
 }: {
   internalState: InternalStateStore;
+  stateRegistry: ProfileStateRegistry;
   tabId: string;
 }): ContextAwarenessToolkit => {
+  const stateAdapters = new Map<string, ProfileStateAdapter<object>>();
+
   return {
     actions: {
       openInNewTab: async (params) => {
@@ -45,5 +56,42 @@ export const createContextAwarenessToolkit = ({
         );
       },
     },
+    getStateAdapter: <TState extends object>(key: ProfileStateKey<TState>) => {
+      if (stateAdapters.has(key.rawKey)) {
+        return stateAdapters.get(key.rawKey) as unknown as ProfileStateAdapter<TState>;
+      }
+
+      stateRegistry.getDefinition(key);
+
+      const getState = () => {
+        const tabState = selectTab(internalState.getState(), tabId);
+        return (tabState?.profileState[key.rawKey] ?? {}) as TState;
+      };
+
+      const state$ = from(internalState).pipe(map(getState), distinctUntilChanged());
+
+      const adapter: ProfileStateAdapter<TState> = {
+        getState,
+        getState$: () => state$,
+        setState: (profileState, _options) => {
+          internalState.dispatch(
+            internalStateActions.setProfileState({ tabId, key: key.rawKey, profileState })
+          );
+        },
+        updateState: (stateUpdate, _options) => {
+          internalState.dispatch(
+            internalStateActions.setProfileState({
+              tabId,
+              key: key.rawKey,
+              profileState: { ...getState(), ...stateUpdate },
+            })
+          );
+        },
+      };
+
+      stateAdapters.set(key.rawKey, adapter as unknown as ProfileStateAdapter<object>);
+
+      return adapter;
+    },
   };
 };

src/platform/plugins/shared/discover/public/application/main/state_management/redux/internal_state.ts

@@ -429,6 +429,11 @@ export const internalStateSlice = createSlice({
         };
       }),
 
+    setProfileState: (state, action: TabAction<{ key: string; profileState: object }>) =>
+      withTab(state, { tabId: state.tabs.unsafeCurrentId }, (tab) => {
+        tab.profileState[action.payload.key] = action.payload.profileState;
+      }),
+
     resetOnSavedSearchChange: (state, action: TabAction) =>
       withTab(state, action.payload, (tab) => {
         tab.overriddenVisContextAfterInvalidation = undefined;
@@ -758,6 +763,7 @@ export const createInternalStateStore = (
     getContextAwarenessToolkit: (tabId: string) => {
       return createContextAwarenessToolkit({
         internalState,
+        stateRegistry: options.services.profileStateRegistry,
         tabId,
       });
     },

src/platform/plugins/shared/discover/public/application/main/state_management/redux/types.ts

@@ -204,6 +204,7 @@ export interface TabState extends TabItem {
   dataRequestParams: InternalStateDataRequestParams;
   overriddenVisContextAfterInvalidation: UnifiedHistogramVisContext | {} | undefined; // it will be used during saving of the Discover Session
   defaultProfileState: DefaultProfileState;
+  profileState: Record<string, object | undefined>;
   uiState: {
     esqlEditor?: Partial<ESQLEditorRestorableState>;
     dataGrid?: Partial<UnifiedDataTableRestorableState>;

src/platform/plugins/shared/discover/public/build_services.ts

@@ -70,7 +70,7 @@ import type { DiscoverStartPlugins } from './types';
 import type { DiscoverContextAppLocator } from './application/context/services/locator';
 import type { DiscoverSingleDocLocator } from './application/doc/locator';
 import type { DiscoverAppLocator } from '../common';
-import type { ProfilesManager } from './context_awareness';
+import { type ProfilesManager, ProfileStateRegistry } from './context_awareness';
 import type { DiscoverEBTManager } from './ebt_manager';
 import {
   CASCADE_LAYOUT_ENABLED_FEATURE_FLAG_KEY,
@@ -157,6 +157,7 @@ export interface DiscoverServices {
   noDataPage?: NoDataPagePluginStart;
   observabilityAIAssistant?: ObservabilityAIAssistantPublicStart;
   profilesManager: ProfilesManager;
+  profileStateRegistry: ProfileStateRegistry;
   ebtManager: DiscoverEBTManager;
   fieldsMetadata?: FieldsMetadataPublicStart;
   logsDataAccess?: LogsDataAccessPluginStart;
@@ -263,6 +264,7 @@ export const buildServices = ({
     noDataPage: plugins.noDataPage,
     observabilityAIAssistant: plugins.observabilityAIAssistant,
     profilesManager,
+    profileStateRegistry: new ProfileStateRegistry(),
     ebtManager,
     fieldsMetadata: plugins.fieldsMetadata,
     logsDataAccess: plugins.logsDataAccess,

src/platform/plugins/shared/discover/public/context_awareness/index.ts

@@ -19,6 +19,14 @@ export {
   type GetProfilesOptions,
 } from './profiles_manager';
 export { type ProfileProviderSharedServices } from './profile_providers';
+export {
+  type ProfileStateAdapter,
+  type ProfileStateDefinition,
+  type ProfileStateDescriptor,
+  ProfileStateKey,
+  ProfileStateType,
+  ProfileStateRegistry,
+} from './profile_state';
 export {
   useProfileAccessor,
   useRootProfile,

src/platform/plugins/shared/discover/public/context_awareness/profile_providers/example/example_data_source_profile/profile.tsx

@@ -7,18 +7,29 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { EuiBadge, EuiFlyout } from '@elastic/eui';
+import {
+  EuiBadge,
+  EuiFlexGroup,
+  EuiFlexItem,
+  EuiFlyout,
+  EuiFormRow,
+  EuiSelect,
+  EuiSpacer,
+  EuiTitle,
+} from '@elastic/eui';
 import type { DataViewField } from '@kbn/data-views-plugin/common';
 import type { RowControlColumn } from '@kbn/discover-utils';
 import { AppMenuActionId, getFieldValue } from '@kbn/discover-utils';
 import { capitalize } from 'lodash';
 import React from 'react';
+import { useObservable } from '@kbn/use-observable';
 import type { DataSourceProfileProvider } from '../../../profiles';
 import { DataSourceCategory } from '../../../profiles';
 import { extractIndexPatternFrom } from '../../extract_index_pattern_from';
 import { ChartWithCustomButtons, CustomDocViewerFooter, CustomDocViewerHeader } from './components';
 import { CustomDocView } from './components/custom_doc_view';
 import { RestorableStateDocView } from './components/restorable_state_doc_view';
+import { COLOR_STATE_KEY } from '../../../profile_state';
 
 export const createExampleDataSourceProfileProvider = (): DataSourceProfileProvider<{
   formatRecord: (flattenedRecord: Record<string, unknown>) => string;
@@ -61,9 +72,11 @@ export const createExampleDataSourceProfileProvider = (): DataSourceProfileProvi
         );
       },
     }),
-    getDocViewer:
-      (prev, { context, toolkit }) =>
-      (params) => {
+    getDocViewer: (prev, { context, toolkit }) => {
+      const stateAdapter = toolkit.getStateAdapter(COLOR_STATE_KEY);
+      const colorState$ = stateAdapter.getState$();
+
+      return (params) => {
         const { openInNewTab, updateESQLQuery } = toolkit.actions;
         const recordId = params.record.id;
         const prevValue = prev(params);
@@ -91,12 +104,51 @@ export const createExampleDataSourceProfileProvider = (): DataSourceProfileProvi
               render: (props) => <RestorableStateDocView {...props} />,
             });
 
+            registry.add({
+              id: 'doc_view_extensible_state_example',
+              title: 'Extensible State Example',
+              order: 2,
+              render: function Render() {
+                const colorState = useObservable(colorState$, stateAdapter.getState());
+                const options = [
+                  { value: 'default', text: 'None' },
+                  { value: 'blue', text: 'Blue' },
+                  { value: 'pink', text: 'Pink' },
+                  { value: 'green', text: 'Green' },
+                ];
+
+                return (
+                  <>
+                    <EuiSpacer size="s" />
+                    <EuiFlexGroup direction="column" gutterSize="m" responsive={false}>
+                      <EuiTitle size="xs">
+                        <h3>Extensible State Example</h3>
+                      </EuiTitle>
+                      <EuiFlexItem grow={false}>
+                        <EuiFormRow label="Favourite color">
+                          <EuiSelect
+                            options={options}
+                            value={colorState.favouriteColor}
+                            onChange={(e) => {
+                              stateAdapter.updateState({ favouriteColor: e.target.value });
+                            }}
+                            aria-label="Select favourite color"
+                          />
+                        </EuiFormRow>
+                      </EuiFlexItem>
+                    </EuiFlexGroup>
+                  </>
+                );
+              },
+            });
+
             return prevValue.docViewsRegistry(registry);
           },
           renderHeader: (props) => <CustomDocViewerHeader {...props} />,
           renderFooter: (props) => <CustomDocViewerFooter {...props} />,
         };
-      },
+      };
+    },
     /**
      * The `getAppMenu` extension point gives access to AppMenuRegistry with methods `registerCustomItem` and
      * `registerCustomPopoverItem`.

src/platform/plugins/shared/discover/public/context_awareness/profile_providers/example/example_root_profile/profile.tsx

@@ -10,27 +10,39 @@
 import { EuiBadge, EuiFlyout } from '@elastic/eui';
 import { getFieldValue } from '@kbn/discover-utils';
 import React from 'react';
+import { useObservable } from '@kbn/use-observable';
 import type { RootProfileProvider } from '../../../profiles';
 import { SolutionType } from '../../../profiles';
+import { COLOR_STATE_KEY } from '../../../profile_state';
 
 export const createExampleRootProfileProvider = (): RootProfileProvider => ({
   profileId: 'example-root-profile',
   isExperimental: true,
   profile: {
     getDefaultAdHocDataViews,
     getDefaultEsqlQuery,
-    getCellRenderers: (prev) => (params) => ({
-      ...prev(params),
-      '@timestamp': (props) => {
-        const timestamp = getFieldValue(props.row, '@timestamp') as string;
+    getCellRenderers: (prev, { toolkit }) => {
+      const stateAdapter = toolkit.getStateAdapter(COLOR_STATE_KEY);
+      const colorState$ = stateAdapter.getState$();
 
-        return (
-          <EuiBadge color="hollow" title={timestamp} data-test-subj="exampleRootProfileTimestamp">
-            {timestamp}
-          </EuiBadge>
-        );
-      },
-    }),
+      return (params) => ({
+        ...prev(params),
+        '@timestamp': function Timestamp(props) {
+          const timestamp = getFieldValue(props.row, '@timestamp') as string;
+          const colorState = useObservable(colorState$, stateAdapter.getState());
+
+          return (
+            <EuiBadge
+              color={colorState.favouriteColor}
+              title={timestamp}
+              data-test-subj="exampleRootProfileTimestamp"
+            >
+              {timestamp}
+            </EuiBadge>
+          );
+        },
+      });
+    },
     /**
      * The `getAppMenu` extension point gives access to AppMenuRegistry with methods `registerCustomItem` and
      * `registerCustomPopoverItem`.

src/platform/plugins/shared/discover/public/context_awareness/profile_state.ts

@@ -0,0 +1,93 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+/* eslint-disable max-classes-per-file */
+
+import type { Observable } from 'rxjs';
+
+export class ProfileStateKey<_TState extends object> {
+  private constructor(public readonly rawKey: string) {}
+
+  static create<TState extends object>(rawKey: string): ProfileStateKey<TState> {
+    return new ProfileStateKey<TState>(rawKey);
+  }
+}
+
+export enum ProfileStateHistoryMethod {
+  Push = 'push',
+  Replace = 'replace',
+}
+
+export interface ProfileStateMutationOptions {
+  historyMethod?: ProfileStateHistoryMethod;
+}
+
+export interface ProfileStateAdapter<TState extends object> {
+  getState: () => TState;
+  getState$: () => Observable<TState>;
+  setState: (state: TState, options?: ProfileStateMutationOptions) => void;
+  updateState: (stateUpdate: Partial<TState>, options?: ProfileStateMutationOptions) => void;
+}
+
+export enum ProfileStateType {
+  Ui = 'ui',
+  Url = 'url',
+  Persistent = 'persistent',
+}
+
+export type ProfileStateDescriptor<TState extends object> = {
+  [key in keyof TState]: {
+    type: ProfileStateType;
+  };
+};
+
+export interface ProfileStateDefinition<TState extends object> {
+  key: ProfileStateKey<TState>;
+  descriptor: ProfileStateDescriptor<TState>;
+}
+
+export class ProfileStateRegistry {
+  private readonly stateDefinitions = new Map<string, ProfileStateDefinition<object>>();
+
+  public registerDefinition<TState extends object>(definition: ProfileStateDefinition<TState>) {
+    if (this.stateDefinitions.has(definition.key.rawKey)) {
+      throw new Error(`State with key ${definition.key.rawKey} is already registered.`);
+    }
+    this.stateDefinitions.set(definition.key.rawKey, definition);
+  }
+
+  public getDefinition<TState extends object>(
+    key: ProfileStateKey<TState>
+  ): ProfileStateDefinition<TState> {
+    const definition = this.stateDefinitions.get(key.rawKey);
+    if (!definition) {
+      throw new Error(`State with key ${key.rawKey} is not registered.`);
+    }
+    return definition as ProfileStateDefinition<TState>;
+  }
+}
+
+export const registerProfileStateDefinitions = (registry: ProfileStateRegistry) => {
+  registry.registerDefinition(colorStateDefinition);
+};
+
+interface ColorState {
+  favouriteColor: string;
+  leastFavouriteColor: string;
+}
+
+export const COLOR_STATE_KEY = ProfileStateKey.create<ColorState>('color_state');
+
+const colorStateDefinition: ProfileStateDefinition<ColorState> = {
+  key: COLOR_STATE_KEY,
+  descriptor: {
+    favouriteColor: { type: ProfileStateType.Ui },
+    leastFavouriteColor: { type: ProfileStateType.Ui },
+  },
+};