Change MethodRow to accept ModelEditorViewState object instead of just Mode

extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx

@@ -235,7 +235,7 @@ export const LibraryRow = ({
             modeledMethods={modeledMethods}
             modifiedSignatures={modifiedSignatures}
             inProgressMethods={inProgressMethods}
-            mode={viewState.mode}
+            viewState={viewState}
             hideModeledMethods={hideModeledMethods}
             revealedMethodSignature={revealedMethodSignature}
             onChange={onChange}

extensions/ql-vscode/src/view/model-editor/MethodRow.tsx

@@ -21,6 +21,7 @@ import { MethodName } from "./MethodName";
 import { ModelTypeDropdown } from "./ModelTypeDropdown";
 import { ModelInputDropdown } from "./ModelInputDropdown";
 import { ModelOutputDropdown } from "./ModelOutputDropdown";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";
 
 const ApiOrMethodCell = styled(VSCodeDataGridCell)`
   display: flex;
@@ -58,7 +59,7 @@ export type MethodRowProps = {
   modeledMethod: ModeledMethod | undefined;
   methodIsUnsaved: boolean;
   modelingInProgress: boolean;
-  mode: Mode;
+  viewState: ModelEditorViewState;
   revealedMethodSignature: string | null;
   onChange: (modeledMethod: ModeledMethod) => void;
 };
@@ -90,7 +91,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
       method,
       modeledMethod,
       methodIsUnsaved,
-      mode,
+      viewState,
       revealedMethodSignature,
       onChange,
     } = props;
@@ -112,7 +113,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
           <ModelingStatusIndicator status={modelingStatus} />
           <MethodClassifications method={method} />
           <MethodName {...props.method} />
-          {mode === Mode.Application && (
+          {viewState.mode === Mode.Application && (
             <UsagesButton onClick={jumpToUsage}>
               {method.usages.length}
             </UsagesButton>
@@ -178,7 +179,7 @@ const UnmodelableMethodRow = forwardRef<
   HTMLElement | undefined,
   MethodRowProps
 >((props, ref) => {
-  const { method, mode, revealedMethodSignature } = props;
+  const { method, viewState, revealedMethodSignature } = props;
 
   const jumpToUsage = useCallback(
     () => sendJumpToUsageMessage(method),
@@ -194,7 +195,7 @@ const UnmodelableMethodRow = forwardRef<
       <ApiOrMethodCell gridColumn={1}>
         <ModelingStatusIndicator status="saved" />
         <MethodName {...props.method} />
-        {mode === Mode.Application && (
+        {viewState.mode === Mode.Application && (
           <UsagesButton onClick={jumpToUsage}>
             {method.usages.length}
           </UsagesButton>

extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx

@@ -8,10 +8,10 @@ import { MethodRow } from "./MethodRow";
 import { Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { useMemo } from "react";
-import { Mode } from "../../model-editor/shared/mode";
 import { sortMethods } from "../../model-editor/shared/sorting";
 import { InProgressMethods } from "../../model-editor/shared/in-progress-methods";
 import { HiddenMethodsRow } from "./HiddenMethodsRow";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";
 
 export const GRID_TEMPLATE_COLUMNS = "0.5fr 0.125fr 0.125fr 0.125fr 0.125fr";
 
@@ -21,7 +21,7 @@ export type ModeledMethodDataGridProps = {
   modeledMethods: Record<string, ModeledMethod>;
   modifiedSignatures: Set<string>;
   inProgressMethods: InProgressMethods;
-  mode: Mode;
+  viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
   onChange: (modeledMethod: ModeledMethod) => void;
@@ -33,7 +33,7 @@ export const ModeledMethodDataGrid = ({
   modeledMethods,
   modifiedSignatures,
   inProgressMethods,
-  mode,
+  viewState,
   hideModeledMethods,
   revealedMethodSignature,
   onChange,
@@ -95,7 +95,7 @@ export const ModeledMethodDataGrid = ({
                 packageName,
                 method.signature,
               )}
-              mode={mode}
+              viewState={viewState}
               revealedMethodSignature={revealedMethodSignature}
               onChange={onChange}
             />

extensions/ql-vscode/src/view/model-editor/__tests__/MethodRow.spec.tsx

@@ -9,6 +9,8 @@ import { Mode } from "../../../model-editor/shared/mode";
 import { MethodRow, MethodRowProps } from "../MethodRow";
 import { ModeledMethod } from "../../../model-editor/modeled-method";
 import userEvent from "@testing-library/user-event";
+import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
 
 describe(MethodRow.name, () => {
   const method = createMethod({
@@ -31,6 +33,14 @@ describe(MethodRow.name, () => {
   };
   const onChange = jest.fn();
 
+  const viewState: ModelEditorViewState = {
+    mode: Mode.Application,
+    showFlowGeneration: false,
+    showLlmButton: false,
+    showMultipleModels: false,
+    extensionPack: createMockExtensionPack(),
+  };
+
   const render = (props: Partial<MethodRowProps> = {}) =>
     reactRender(
       <MethodRow
@@ -40,7 +50,7 @@ describe(MethodRow.name, () => {
         methodIsUnsaved={false}
         modelingInProgress={false}
         revealedMethodSignature={null}
-        mode={Mode.Application}
+        viewState={viewState}
         onChange={onChange}
         {...props}
       />,

extensions/ql-vscode/src/view/model-editor/__tests__/ModeledMethodDataGrid.spec.tsx

@@ -7,6 +7,8 @@ import {
   ModeledMethodDataGrid,
   ModeledMethodDataGridProps,
 } from "../ModeledMethodDataGrid";
+import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
 
 describe(ModeledMethodDataGrid.name, () => {
   const method1 = createMethod({
@@ -41,6 +43,14 @@ describe(ModeledMethodDataGrid.name, () => {
   });
   const onChange = jest.fn();
 
+  const viewState: ModelEditorViewState = {
+    mode: Mode.Application,
+    showFlowGeneration: false,
+    showLlmButton: false,
+    showMultipleModels: false,
+    extensionPack: createMockExtensionPack(),
+  };
+
   const render = (props: Partial<ModeledMethodDataGridProps> = {}) =>
     reactRender(
       <ModeledMethodDataGrid
@@ -58,7 +68,7 @@ describe(ModeledMethodDataGrid.name, () => {
         }}
         modifiedSignatures={new Set([method1.signature])}
         inProgressMethods={new InProgressMethods()}
-        mode={Mode.Application}
+        viewState={viewState}
         hideModeledMethods={false}
         revealedMethodSignature={null}
         onChange={onChange}