diff --git a/packages/editor-ui/src/components/CodeNodeEditor/CodeNodeEditor.vue b/packages/editor-ui/src/components/CodeNodeEditor/CodeNodeEditor.vue
index 1d98ae32c1..9deed663d1 100644
--- a/packages/editor-ui/src/components/CodeNodeEditor/CodeNodeEditor.vue
+++ b/packages/editor-ui/src/components/CodeNodeEditor/CodeNodeEditor.vue
@@ -7,7 +7,7 @@
 	>
 		<div ref="codeNodeEditor" class="code-node-editor-input ph-no-capture"></div>
 		<n8n-button
-			v-if="isCloud && (isEditorHovered || isEditorFocused)"
+			v-if="aiButtonEnabled && (isEditorHovered || isEditorFocused)"
 			size="small"
 			type="tertiary"
 			:class="$style['ask-ai-button']"
@@ -38,7 +38,6 @@ import { workflowHelpers } from '@/mixins/workflowHelpers'; // for json field co
 import { ASK_AI_MODAL_KEY, CODE_NODE_TYPE } from '@/constants';
 import { codeNodeEditorEventBus } from '@/event-bus';
 import { useRootStore } from '@/stores/n8nRoot.store';
-import { useSettingsStore } from '@/stores/settings.store';
 
 import { readOnlyEditorExtensions, writableEditorExtensions } from './baseExtensions';
 import { CODE_PLACEHOLDERS } from './constants';
@@ -50,6 +49,10 @@ export default defineComponent({
 	name: 'code-node-editor',
 	mixins: [linterExtension, completerExtension, workflowHelpers],
 	props: {
+		aiButtonEnabled: {
+			type: Boolean,
+			default: false,
+		},
 		mode: {
 			type: String as PropType<CodeExecutionMode>,
 			validator: (value: CodeExecutionMode): boolean => CODE_EXECUTION_MODES.includes(value),
@@ -97,9 +100,6 @@ export default defineComponent({
 	},
 	computed: {
 		...mapStores(useRootStore),
-		isCloud() {
-			return useSettingsStore().deploymentType === 'cloud';
-		},
 		content(): string {
 			if (!this.editor) return '';
 
diff --git a/packages/editor-ui/src/components/ParameterInput.vue b/packages/editor-ui/src/components/ParameterInput.vue
index 244570a4d3..682c465dba 100644
--- a/packages/editor-ui/src/components/ParameterInput.vue
+++ b/packages/editor-ui/src/components/ParameterInput.vue
@@ -91,6 +91,7 @@
 					:defaultValue="parameter.default"
 					:language="editorLanguage"
 					:isReadOnly="isReadOnly"
+					:aiButtonEnabled="settingsStore.isCloudDeployment"
 					@valueChanged="valueChangedDebounced"
 				/>
 
@@ -387,6 +388,7 @@ import { useWorkflowsStore } from '@/stores/workflows.store';
 import { useNDVStore } from '@/stores/ndv.store';
 import { useNodeTypesStore } from '@/stores/nodeTypes.store';
 import { useCredentialsStore } from '@/stores/credentials.store';
+import { useSettingsStore } from '@/stores/settings.store';
 import { htmlEditorEventBus } from '@/event-bus';
 import Vue from 'vue';
 
@@ -519,7 +521,13 @@ export default defineComponent({
 		},
 	},
 	computed: {
-		...mapStores(useCredentialsStore, useNodeTypesStore, useNDVStore, useWorkflowsStore),
+		...mapStores(
+			useCredentialsStore,
+			useNodeTypesStore,
+			useNDVStore,
+			useWorkflowsStore,
+			useSettingsStore,
+		),
 		expressionDisplayValue(): string {
 			if (this.forceShowExpression) {
 				return '';