From 958cb6d70e4cda715309fa3eb06dfba7e42645ef Mon Sep 17 00:00:00 2001
From: Yohann Paris <github@yohannparis.com>
Date: Wed, 2 Oct 2024 16:36:40 -0400
Subject: [PATCH] Update tera-model in operators (#4987)

---
 .../model/petrinet/tera-model-description.vue |  2 +-
 .../src/components/model/tera-model.vue       |  3 +-
 .../components/workflow/ops/model-edit/mod.ts |  2 +-
 ...edit.vue => tera-model-edit-drilldown.vue} | 76 ++++-------------
 .../tera-model-from-equations-drilldown.vue   | 71 ++++++----------
 .../workflow/ops/stratify-mira/mod.ts         |  4 +-
 ...y-mira.vue => tera-stratify-drilldown.vue} | 85 +++++++------------
 ...y-node-mira.vue => tera-stratify-node.vue} |  0
 .../hmi-client/src/composables/project.ts     |  3 +-
 .../hmi-client/src/page/WorkflowNode.vue      |  2 +-
 10 files changed, 80 insertions(+), 168 deletions(-)
 rename packages/client/hmi-client/src/components/workflow/ops/model-edit/{tera-model-edit.vue => tera-model-edit-drilldown.vue} (86%)
 rename packages/client/hmi-client/src/components/workflow/ops/stratify-mira/{tera-stratify-mira.vue => tera-stratify-drilldown.vue} (87%)
 rename packages/client/hmi-client/src/components/workflow/ops/stratify-mira/{tera-stratify-node-mira.vue => tera-stratify-node.vue} (100%)

diff --git a/packages/client/hmi-client/src/components/model/petrinet/tera-model-description.vue b/packages/client/hmi-client/src/components/model/petrinet/tera-model-description.vue
index 530e78ffc3..b6ace60312 100644
--- a/packages/client/hmi-client/src/components/model/petrinet/tera-model-description.vue
+++ b/packages/client/hmi-client/src/components/model/petrinet/tera-model-description.vue
@@ -48,7 +48,7 @@ const props = defineProps<{
 const emit = defineEmits(['update-model']);
 const teraModelDiagramRef = ref();
 
-const currentActiveIndexes = ref([0, 1, 2, 3]);
+const currentActiveIndexes = ref([1, 2, 3]);
 const relatedTerariumArtifacts = ref<Asset[]>([]);
 const relatedTerariumModels = computed(() => relatedTerariumArtifacts.value.filter((d) => isModel(d)) as Model[]);
 const relatedTerariumDatasets = computed(() => relatedTerariumArtifacts.value.filter((d) => isDataset(d)) as Dataset[]);
diff --git a/packages/client/hmi-client/src/components/model/tera-model.vue b/packages/client/hmi-client/src/components/model/tera-model.vue
index bfd97ea92d..73ae5b166b 100644
--- a/packages/client/hmi-client/src/components/model/tera-model.vue
+++ b/packages/client/hmi-client/src/components/model/tera-model.vue
@@ -42,6 +42,7 @@
 					:disabled="!(hasChanged && hasEditPermission)"
 				/>
 				<Button
+					v-if="!isWorkflow"
 					:label="`Save ${isWorkflow ? 'for re-use' : 'as'}`"
 					severity="secondary"
 					outlined
@@ -70,7 +71,7 @@
 		:asset-type="AssetType.Model"
 		:initial-name="temporaryModel?.header.name"
 		:is-visible="showSaveModal"
-		:open-on-save="isWorkflow ? false : true"
+		:open-on-save="!isWorkflow"
 		@close-modal="showSaveModal = false"
 		@on-save="onModalSave"
 	/>
diff --git a/packages/client/hmi-client/src/components/workflow/ops/model-edit/mod.ts b/packages/client/hmi-client/src/components/workflow/ops/model-edit/mod.ts
index ff5e492113..54b44bcf43 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/model-edit/mod.ts
+++ b/packages/client/hmi-client/src/components/workflow/ops/model-edit/mod.ts
@@ -1,5 +1,5 @@
 import node from './tera-model-edit-node.vue';
-import drilldown from './tera-model-edit.vue';
+import drilldown from './tera-model-edit-drilldown.vue';
 import { ModelEditOperation as operation } from './model-edit-operation';
 
 const name = operation.name;
diff --git a/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit.vue b/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit-drilldown.vue
similarity index 86%
rename from packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit.vue
rename to packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit-drilldown.vue
index ef5068b8ca..49add663f1 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit-drilldown.vue
@@ -50,83 +50,47 @@
 				:options="outputs"
 				is-selectable
 			>
-				<section class="right-side">
-					<Button
-						class="mr-3"
-						outlined
-						severity="secondary"
-						:disabled="isSaveDisabled"
-						label="Save for reuse"
-						@click="showSaveModelModal = true"
-					/>
-				</section>
 				<tera-notebook-error
 					v-if="executeResponse.status === OperatorStatus.ERROR"
 					:name="executeResponse.name"
 					:value="executeResponse.value"
 					:traceback="executeResponse.traceback"
 				/>
-				<template v-else-if="amr">
-					<tera-model-diagram :model="amr" />
-					<tera-model-parts :model="amr" :feature-config="{ isPreview: true }" />
-				</template>
+				<tera-model v-else-if="amr" is-workflow :assetId="amr.id" @on-save="updateNode" />
 				<tera-progress-spinner v-else-if="isUpdatingModel || !amr" is-centered :font-size="2">
 					Loading...
 				</tera-progress-spinner>
 			</tera-drilldown-preview>
 		</template>
 		<tera-drilldown-section :tabName="DrilldownTabs.Wizard">
-			<!-- <tera-model-template-editor v-if="amr"
-					:model="amr"
-					:kernel-manager="kernelManager"
-					@output-code="(data: any) => appendCode(data, 'executed_code')"
-					@sync-with-mira-model="syncWithMiraModel"
-					@save-new-model-output="createOutput"
-					@reset="resetModel"
-				/> -->
 			<p class="m-4">Wizard is disabled for now.</p>
 		</tera-drilldown-section>
 	</tera-drilldown>
-	<tera-save-asset-modal
-		v-if="amr"
-		:asset="amr"
-		:initial-name="amr.name"
-		:assetType="AssetType.Model"
-		:is-updating-asset="true"
-		:is-visible="showSaveModelModal"
-		@close-modal="showSaveModelModal = false"
-		@on-save="updateNode"
-	/>
 </template>
 
 <script setup lang="ts">
 import { cloneDeep, isEmpty, isEqual, debounce } from 'lodash';
 import { computed, onMounted, onUnmounted, ref, watch } from 'vue';
-import Button from 'primevue/button';
+import '@/ace-config';
 import { VAceEditor } from 'vue3-ace-editor';
 import { VAceEditorInstance } from 'vue3-ace-editor/types';
-import '@/ace-config';
+import Button from 'primevue/button';
+import { DrilldownTabs } from '@/types/common';
 import type { Model } from '@/types/Types';
-import { AssetType } from '@/types/Types';
-import { createModel, getModel } from '@/services/model';
 import { OperatorStatus, WorkflowNode, WorkflowOutput } from '@/types/workflow';
+import { KernelSessionManager } from '@/services/jupyter';
+import { createModel, getModel } from '@/services/model';
+import { getModelIdFromModelConfigurationId } from '@/services/model-configurations';
+import { saveCodeToState } from '@/services/notebook';
 import { logger } from '@/utils/logger';
-import TeraModelDiagram from '@/components/model/petrinet/model-diagrams/tera-model-diagram.vue';
-import TeraModelParts from '@/components/model/tera-model-parts.vue';
-import TeraNotebookError from '@/components/drilldown/tera-notebook-error.vue';
 import TeraDrilldown from '@/components/drilldown/tera-drilldown.vue';
 import TeraDrilldownPreview from '@/components/drilldown/tera-drilldown-preview.vue';
 import TeraDrilldownSection from '@/components/drilldown/tera-drilldown-section.vue';
-import TeraProgressSpinner from '@/components/widgets/tera-progress-spinner.vue';
-// import TeraModelTemplateEditor from '@/components/model-template/tera-model-template-editor.vue';
+import TeraModel from '@/components/model/tera-model.vue';
+import TeraNotebookError from '@/components/drilldown/tera-notebook-error.vue';
 import TeraNotebookJupyterInput from '@/components/llm/tera-notebook-jupyter-input.vue';
-import { KernelSessionManager } from '@/services/jupyter';
-import { getModelIdFromModelConfigurationId } from '@/services/model-configurations';
-import TeraSaveAssetModal from '@/components/project/tera-save-asset-modal.vue';
-import { saveCodeToState } from '@/services/notebook';
-import { DrilldownTabs } from '@/types/common';
+import TeraProgressSpinner from '@/components/widgets/tera-progress-spinner.vue';
 import { nodeOutputLabel } from '@/components/workflow/util';
-import { useProjects } from '@/composables/project';
 import { ModelEditOperationState, ModelEditOperation } from './model-edit-operation';
 
 const props = defineProps<{
@@ -158,7 +122,6 @@ const isUpdatingModel = ref(false);
 const kernelManager = new KernelSessionManager();
 const amr = ref<Model | null>(null);
 let activeModelId: string | null = null;
-const showSaveModelModal = ref(false);
 
 let editor: VAceEditorInstance['_editor'] | null;
 const sampleAgentQuestions = [
@@ -341,7 +304,7 @@ const handleOutputChange = async () => {
 		const jupyterContext = buildJupyterContext();
 		if (jupyterContext) {
 			if (kernelManager.jupyterSession !== null) {
-				// when coming from output dropdown change we should shutdown first
+				// when coming from output dropdown change we should shut down first
 				kernelManager.shutdown();
 			}
 			await kernelManager.init('beaker_kernel', 'Beaker Kernel', jupyterContext);
@@ -387,13 +350,6 @@ watch(
 	{ immediate: true }
 );
 
-const isSaveDisabled = computed(() => {
-	const id = amr.value?.id;
-	if (!id) return true;
-	const outputPort = props.node.outputs?.find((port) => port.value?.[0] === id);
-	return useProjects().hasAssetInActiveProject(outputPort?.value?.[0]);
-});
-
 function updateNode(model: Model) {
 	const id = amr.value?.id;
 	if (!id || !model) return;
@@ -408,7 +364,7 @@ function updateNode(model: Model) {
 }
 
 onMounted(async () => {
-	// By default the first output option is the original model
+	// By default, the first output option is the original model
 	if (isReadyToCreateDefaultOutput.value) {
 		const input = props.node.inputs[0];
 		if (!input) return;
@@ -452,15 +408,15 @@ onUnmounted(() => {
 }
 
 .input-small {
-	padding: 0.5rem;
+	padding: var(--gap-2);
 	width: 100%;
 }
 
 .code-executed-warning {
 	background-color: #ffe6e6;
+	border-radius: var(--border-radius);
 	color: #cc0000;
-	padding: 10px;
-	border-radius: 4px;
+	padding: var(--gap-2-5);
 }
 
 .right-side {
diff --git a/packages/client/hmi-client/src/components/workflow/ops/model-from-equations/tera-model-from-equations-drilldown.vue b/packages/client/hmi-client/src/components/workflow/ops/model-from-equations/tera-model-from-equations-drilldown.vue
index f1aad80bc9..364c10c1fb 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/model-from-equations/tera-model-from-equations-drilldown.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/model-from-equations/tera-model-from-equations-drilldown.vue
@@ -119,80 +119,53 @@
 			</tera-slider-panel>
 			<tera-slider-panel v-model:is-open="isOutputOpen" header="Output" content-width="100%">
 				<template #content>
-					<header class="flex align-items-center p-3">
-						<h4>Equation conversions</h4>
-						<Button
-							v-if="selectedModel"
-							label="Save for re-use"
-							outlined
-							severity="secondary"
-							class="ml-auto"
-							@click="showSaveModelModal = true"
-						/>
-					</header>
 					<tera-drilldown-preview>
-						<tera-model-description v-if="selectedModel" :model="selectedModel" :generating-card="isGeneratingCard" />
-						<tera-operator-placeholder v-else :node="node" style="height: 100%" />
+						<tera-model v-if="selectedModel" is-workflow :assetId="selectedModel.id" @on-save="onModelSaveEvent" />
+						<tera-operator-placeholder v-else :node="node" class="h-100" />
 					</tera-drilldown-preview>
 				</template>
 			</tera-slider-panel>
 		</template>
 	</tera-drilldown>
-	<tera-save-asset-modal
-		v-if="selectedModel"
-		:asset="selectedModel"
-		:is-visible="showSaveModelModal"
-		@close-modal="onCloseModelModal"
-	/>
 </template>
 
 <script setup lang="ts">
 import { AssetBlock, WorkflowNode } from '@/types/workflow';
+import Button from 'primevue/button';
 import Checkbox from 'primevue/checkbox';
+import Textarea from 'primevue/textarea';
+import TeraAssetBlock from '@/components/widgets/tera-asset-block.vue';
 import TeraDrilldown from '@/components/drilldown/tera-drilldown.vue';
 import TeraDrilldownPreview from '@/components/drilldown/tera-drilldown-preview.vue';
-import TeraAssetBlock from '@/components/widgets/tera-asset-block.vue';
 import { computed, onMounted, ref, watch } from 'vue';
-import { downloadDocumentAsset, getDocumentAsset, getDocumentFileAsText } from '@/services/document-assets';
 import type { Card, DocumentAsset, Model } from '@/types/Types';
 import { cloneDeep, isEmpty } from 'lodash';
 import { equationsToAMR, type EquationsToAMRRequest } from '@/services/knowledge';
-import Button from 'primevue/button';
+import { downloadDocumentAsset, getDocumentAsset, getDocumentFileAsText } from '@/services/document-assets';
+import { enrichModelMetadata } from '@/services/goLLM';
 import { getModel, updateModel } from '@/services/model';
-import TeraOperatorPlaceholder from '@/components/operator/tera-operator-placeholder.vue';
 import { useProjects } from '@/composables/project';
-import TeraMathEditor from '@/components/mathml/tera-math-editor.vue';
-import Textarea from 'primevue/textarea';
+import TeraOperatorPlaceholder from '@/components/operator/tera-operator-placeholder.vue';
+import TeraModel from '@/components/model/tera-model.vue';
 import TeraInputText from '@/components/widgets/tera-input-text.vue';
-import TeraSaveAssetModal from '@/components/project/tera-save-asset-modal.vue';
-import TeraModelDescription from '@/components/model/petrinet/tera-model-description.vue';
-import { enrichModelMetadata } from '@/services/goLLM';
+import TeraMathEditor from '@/components/mathml/tera-math-editor.vue';
 import TeraSliderPanel from '@/components/widgets/tera-slider-panel.vue';
-
+import TeraDrilldownSection from '@/components/drilldown/tera-drilldown-section.vue';
 import TeraPdfEmbed from '@/components/widgets/tera-pdf-embed.vue';
 import TeraTextEditor from '@/components/documents/tera-text-editor.vue';
-import TeraDrilldownSection from '@/components/drilldown/tera-drilldown-section.vue';
-import { ModelFromEquationsState, EquationBlock } from './model-from-equations-operation';
+import { ModelFromEquationsOperation, ModelFromEquationsState, EquationBlock } from './model-from-equations-operation';
 
 const emit = defineEmits(['close', 'update-state', 'append-output', 'select-output']);
 const props = defineProps<{
 	node: WorkflowNode<ModelFromEquationsState>;
 }>();
 
-enum ModelFramework {
-	PetriNet = 'petrinet',
-	RegNet = 'regnet',
-	Decapode = 'decapode',
-	GeneralizedAMR = 'gamr',
-	MathExpressionTree = 'met'
-}
-
 const selectedOutputId = ref<string>('');
 
 const clonedState = ref<ModelFromEquationsState>({
 	equations: [],
 	text: '',
-	modelFramework: ModelFramework.PetriNet,
+	modelFramework: 'petrinet',
 	modelId: null
 });
 
@@ -228,7 +201,6 @@ const selectedModel = ref<Model | null>(null);
 const card = ref<Card | null>(null);
 const goLLMCard = computed<any>(() => document.value?.metadata?.gollmCard);
 
-const showSaveModelModal = ref(false);
 const isGeneratingCard = ref(false);
 const multipleEquations = ref<string>('');
 
@@ -361,10 +333,6 @@ async function fetchModel() {
 	loadingModel.value = false;
 }
 
-function onCloseModelModal() {
-	showSaveModelModal.value = false;
-}
-
 function getEquations() {
 	const newEquations = multipleEquations.value.split('\n');
 	newEquations.forEach((equation) => {
@@ -385,6 +353,19 @@ function getEquationErrorLabel(equation) {
 	return equation.asset.extractionError ? "Couldn't extract equation" : '';
 }
 
+function onModelSaveEvent(event: any) {
+	const state = cloneDeep(props.node.state);
+	state.modelId = event.id;
+	emit('update-state', state);
+	emit('append-output', {
+		type: ModelFromEquationsOperation.outputs[0].type,
+		label: event.header.name,
+		value: [event.id],
+		state,
+		isSelected: false
+	});
+}
+
 // generates the model card and fetches the model when finished
 async function generateCard(modelId: string, docId: string) {
 	isGeneratingCard.value = true;
diff --git a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/mod.ts b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/mod.ts
index d4604e97a9..3e345458bc 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/mod.ts
+++ b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/mod.ts
@@ -1,6 +1,6 @@
 import { StratifyMiraOperation as operation } from './stratify-mira-operation';
-import node from './tera-stratify-node-mira.vue';
-import drilldown from './tera-stratify-mira.vue';
+import node from './tera-stratify-node.vue';
+import drilldown from './tera-stratify-drilldown.vue';
 
 const name = operation.name;
 
diff --git a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-mira.vue b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-drilldown.vue
similarity index 87%
rename from packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-mira.vue
rename to packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-drilldown.vue
index 4374ed391f..665868ea2b 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-mira.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-drilldown.vue
@@ -75,25 +75,13 @@
 				v-model:output="selectedOutputId"
 				is-selectable
 			>
-				<Button
-					class="ml-auto py-3"
-					label="Save for re-use"
-					size="small"
-					outlined
-					:disabled="isSaveDisabled"
-					severity="secondary"
-					@click="showSaveModelModal = true"
-				/>
 				<tera-notebook-error
 					v-if="executeResponse.status === OperatorStatus.ERROR"
 					:name="executeResponse.name"
 					:value="executeResponse.value"
 					:traceback="executeResponse.traceback"
 				/>
-				<template v-else-if="outputAmr">
-					<tera-model-diagram :model="outputAmr" />
-					<tera-model-parts :model="outputAmr" :feature-config="{ isPreview: true }" />
-				</template>
+				<tera-model v-else-if="outputAmr" is-workflow :assetId="outputAmr.id" @on-save="updateNode" />
 				<template v-else>
 					<tera-progress-spinner v-if="isStratifyInProgress" is-centered :font-size="2">
 						Processing...
@@ -103,33 +91,18 @@
 			</tera-drilldown-preview>
 		</template>
 	</tera-drilldown>
-	<tera-save-asset-modal
-		v-if="outputAmr"
-		:asset="outputAmr"
-		:assetType="AssetType.Model"
-		:initial-name="outputAmr.name"
-		:is-visible="showSaveModelModal"
-		:is-updating-asset="true"
-		@close-modal="showSaveModelModal = false"
-	/>
 </template>
 
 <script setup lang="ts">
-import { computed, onMounted, onUnmounted, ref, watch } from 'vue';
-import _ from 'lodash';
+import { onMounted, onUnmounted, ref, watch } from 'vue';
+import { cloneDeep, debounce, isEqual, last } from 'lodash';
 import TeraDrilldownPreview from '@/components/drilldown/tera-drilldown-preview.vue';
 import TeraDrilldownSection from '@/components/drilldown/tera-drilldown-section.vue';
 import TeraDrilldown from '@/components/drilldown/tera-drilldown.vue';
-import TeraModelDiagram from '@/components/model/petrinet/model-diagrams/tera-model-diagram.vue';
 import TeraOperatorPlaceholder from '@/components/operator/tera-operator-placeholder.vue';
-import TeraModelParts from '@/components/model/tera-model-parts.vue';
-import TeraSaveAssetModal from '@/components/project/tera-save-asset-modal.vue';
 import TeraStratificationGroupForm from '@/components/workflow/ops/stratify-mira/tera-stratification-group-form.vue';
 import TeraNotebookJupyterInput from '@/components/llm/tera-notebook-jupyter-input.vue';
-
 import { createModel, getModel } from '@/services/model';
-import { useProjects } from '@/composables/project';
-
 import { WorkflowNode, OperatorStatus } from '@/types/workflow';
 import { logger } from '@/utils/logger';
 import Button from 'primevue/button';
@@ -139,19 +112,17 @@ import { VAceEditorInstance } from 'vue3-ace-editor/types';
 import '@/ace-config';
 import TeraNotebookError from '@/components/drilldown/tera-notebook-error.vue';
 import type { Model } from '@/types/Types';
-import { AssetType } from '@/types/Types';
 import { AMRSchemaNames } from '@/types/common';
 import { getModelIdFromModelConfigurationId } from '@/services/model-configurations';
 import TeraProgressSpinner from '@/components/widgets/tera-progress-spinner.vue';
-
-/* Jupyter imports */
 import { KernelSessionManager } from '@/services/jupyter';
+import TeraModel from '@/components/model/tera-model.vue';
 import { blankStratifyGroup, StratifyGroup, StratifyOperationStateMira } from './stratify-mira-operation';
 
 const props = defineProps<{
 	node: WorkflowNode<StratifyOperationStateMira>;
 }>();
-const emit = defineEmits(['append-output', 'update-state', 'close', 'select-output']);
+const emit = defineEmits(['append-output', 'update-state', 'close', 'select-output', 'update-output']);
 
 enum StratifyTabs {
 	Wizard = 'Wizard',
@@ -167,7 +138,6 @@ const executeResponse = ref({
 	traceback: ''
 });
 const modelNodeOptions = ref<string[]>([]);
-const showSaveModelModal = ref(false);
 
 const isDraft = ref(false);
 
@@ -194,7 +164,7 @@ const updateLlmQuery = (query: string) => {
 };
 
 const updateStratifyGroupForm = (config: StratifyGroup) => {
-	const state = _.cloneDeep(props.node.state);
+	const state = cloneDeep(props.node.state);
 	state.strataGroup = config;
 	emit('update-state', state);
 };
@@ -288,8 +258,8 @@ const handleModelPreview = async (data: any) => {
 		label: newName,
 		type: 'modelId',
 		state: {
-			strataGroup: _.cloneDeep(props.node.state.strataGroup),
-			strataCodeHistory: _.cloneDeep(props.node.state.strataCodeHistory)
+			strataGroup: cloneDeep(props.node.state.strataGroup),
+			strataCodeHistory: cloneDeep(props.node.state.strataCodeHistory)
 		},
 		value: [modelData.id]
 	});
@@ -430,7 +400,7 @@ const runCodeStratify = () => {
 
 // FIXME: Copy pasted in 3 locations, could be written cleaner and in a service. Migrate it to use saveCodeToState from @/services/notebook
 const saveCodeToState = (code: string, hasCodeBeenRun: boolean) => {
-	const state = _.cloneDeep(props.node.state);
+	const state = cloneDeep(props.node.state);
 	state.hasCodeBeenRun = hasCodeBeenRun;
 	// for now only save the last code executed, may want to save all code executed in the future
 	const codeHistoryLength = props.node.state.strataCodeHistory.length;
@@ -448,23 +418,28 @@ const onSelection = (id: string) => {
 	emit('select-output', id);
 };
 
-const isSaveDisabled = computed(() => {
-	const id = amr.value?.id;
-	if (!id || _.isEmpty(selectedOutputId.value)) return true;
-	const outputPort = props.node.outputs?.find((port) => port.id === selectedOutputId.value);
-
-	return useProjects().hasAssetInActiveProject(outputPort?.value?.[0]);
-});
-
 // check if user has made changes to the code
 const hasCodeChange = () => {
 	if (props.node.state.strataCodeHistory.length) {
-		isDraft.value = !_.isEqual(codeText.value, props.node.state.strataCodeHistory?.[0]?.code);
+		isDraft.value = !isEqual(codeText.value, props.node.state.strataCodeHistory?.[0]?.code);
 	} else {
-		isDraft.value = !_.isEqual(codeText.value, '');
+		isDraft.value = !isEqual(codeText.value, '');
 	}
 };
-const checkForCodeChange = _.debounce(hasCodeChange, 100);
+const checkForCodeChange = debounce(hasCodeChange, 100);
+
+function updateNode(model: Model) {
+	const id = outputAmr.value?.id;
+	if (!id || !model) return;
+
+	outputAmr.value = model;
+	const outputPort = cloneDeep(props.node.outputs?.find((port) => port.value?.[0] === id));
+
+	if (!outputPort) return;
+	outputPort.label = model.header.name;
+
+	emit('update-output', outputPort);
+}
 
 watch(
 	() => codeText.value,
@@ -484,7 +459,7 @@ watch(
 			}
 			const modelIdToLoad = output.value?.[0];
 			outputAmr.value = await getModel(modelIdToLoad);
-			codeText.value = _.last(props.node.state.strataCodeHistory)?.code ?? '';
+			codeText.value = last(props.node.state.strataCodeHistory)?.code ?? '';
 		}
 	},
 	{ immediate: true }
@@ -518,10 +493,10 @@ onUnmounted(() => {
 }
 
 .code-executed-warning {
-	background-color: #ffe6e6;
-	color: #cc0000;
-	padding: 10px;
-	border-radius: 4px;
+	background-color: var(--error-message-background);
+	border-radius: var(--border-radius);
+	color: var(--error-message-color);
+	padding: var(--gap-2-5);
 }
 
 .form-section {
diff --git a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-node-mira.vue b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-node.vue
similarity index 100%
rename from packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-node-mira.vue
rename to packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-node.vue
diff --git a/packages/client/hmi-client/src/composables/project.ts b/packages/client/hmi-client/src/composables/project.ts
index 0fd70c7cff..80a9a92dee 100644
--- a/packages/client/hmi-client/src/composables/project.ts
+++ b/packages/client/hmi-client/src/composables/project.ts
@@ -110,8 +110,7 @@ export function useProjects() {
 	 * @returns {ProjectAsset['assetName']}
 	 */
 	function getAssetName(assetId: ProjectAsset['assetId']): ProjectAsset['assetName'] {
-		const asset = findAsset(assetId);
-		return asset?.assetName ?? '';
+		return findAsset(assetId)?.assetName ?? '';
 	}
 
 	/**
diff --git a/packages/client/hmi-client/src/page/WorkflowNode.vue b/packages/client/hmi-client/src/page/WorkflowNode.vue
index 9a90ddb1ba..16f716f990 100644
--- a/packages/client/hmi-client/src/page/WorkflowNode.vue
+++ b/packages/client/hmi-client/src/page/WorkflowNode.vue
@@ -34,7 +34,7 @@ import TeraCalibrateEnsembleCiemssDrilldown from '@/components/workflow/ops/cali
 import TeraSimulateCiemssDrilldown from '@/components/workflow/ops/simulate-ciemss/tera-simulate-ciemss-drilldown.vue';
 import TeraSimulateEnsembleCiemssDrilldown from '@/components/workflow/ops/simulate-ensemble-ciemss/tera-simulate-ensemble-ciemss-drilldown.vue';
 import TeraFunman from '@/components/workflow/ops/funman/tera-funman-drilldown.vue';
-import TeraStratifyMira from '@/components/workflow/ops/stratify-mira/tera-stratify-mira.vue';
+import TeraStratifyMira from '@/components/workflow/ops/stratify-mira/tera-stratify-drilldown.vue';
 import TeraCodeAssetWrapper from '@/components/workflow/ops/code-asset/tera-code-asset-wrapper.vue';
 
 const props = defineProps<{ nodeId: string; workflowId: string }>();