diff --git a/src/app/shared/custom-connection/custom-connection.ts b/src/app/shared/custom-connection/custom-connection.ts
index 1375818..3c91330 100644
--- a/src/app/shared/custom-connection/custom-connection.ts
+++ b/src/app/shared/custom-connection/custom-connection.ts
@@ -1,4 +1,4 @@
-import { ChangeDetectionStrategy, Component, inject, Input } from '@angular/core';
+import { Component, inject, Input } from '@angular/core';
 import { ClassicPreset } from 'rete';
 import { FLOW_DEPENDANT_PORT_KEY, FLOW_DEPENDENCY_PORT_KEY } from '@models/flow';
 import { GraphSelectionService } from '@services/graph-selection/graph-selection';
@@ -79,8 +79,7 @@ import { GraphSelectionService } from '@services/graph-selection/graph-selection
       overflow: visible;
       pointer-events: auto;
     }
-  `],
-  changeDetection: ChangeDetectionStrategy.OnPush
+  `]
 })
 export class CustomConnectionComponent {
   private readonly graphSelection = inject(GraphSelectionService);
diff --git a/src/app/shared/nodes/container-node/container-node.html b/src/app/shared/nodes/container-node/container-node.html
index 534bd1e..8b46fb0 100644
--- a/src/app/shared/nodes/container-node/container-node.html
+++ b/src/app/shared/nodes/container-node/container-node.html
@@ -12,7 +12,11 @@
   }
   <div class="container-node__header" [class.cursor-move]="!isReadonly">
     <div class="container-node__icon">
-      <i class="bi bi-box-seam"></i>
+      @if (nodeIcon().type === 'class') {
+      <i [class]="nodeIcon().value"></i>
+      } @else {
+      <img [src]="nodeIcon().value" alt="Container" class="container-node__icon-image" />
+      }
     </div>
     <div class="container-node__titles">
       <div class="container-node__eyebrow">{{ containerLabel }}</div>
diff --git a/src/app/shared/nodes/container-node/container-node.ts b/src/app/shared/nodes/container-node/container-node.ts
index ebedc43..c91c005 100644
--- a/src/app/shared/nodes/container-node/container-node.ts
+++ b/src/app/shared/nodes/container-node/container-node.ts
@@ -13,7 +13,7 @@ import { EditorStateHolder } from '@stores/flow-editor';
 import { CONTAINER_SUBFLOW_DRAG_MIME } from './container-node-drag';
 import { firstValueFrom } from 'rxjs';
 import { extractSchemaRequirements, SchemaRequirements } from '../schema-requirements';
-import { evaluateUiConditionRule, getValueByPath, parentPath, pathToLabel, readEffectiveUiVisibleConditionRule, readUiConditionRule, resolveSchemaPath, resolveSchemaRef, schemaFieldLabel, shouldSkipSchemaField, splitTemplatedTextParts, valueToDisplayString } from '../node-utility';
+import { evaluateUiConditionRule, getValueByPath, parentPath, pathToLabel, readEffectiveUiVisibleConditionRule, readUiConditionRule, resolveNodeIcon, resolveSchemaPath, resolveSchemaRef, schemaFieldLabel, shouldSkipSchemaField, splitTemplatedTextParts, valueToDisplayString } from '../node-utility';
 
 type ContainerFieldType = 'string' | 'number' | 'integer' | 'boolean' | 'unknown';
 
@@ -142,6 +142,16 @@ export class ContainerNodeComponent {
     return pathToLabel(this.typeName.replace(/Container$/, ' Container'));
   }
 
+  nodeIcon(): { type: 'class' | 'img'; value: string } {
+    const icon = resolveNodeIcon(this.containerSchema, false);
+    if (icon.type === 'class' && icon.value === 'bi bi-person-check-fill') {
+      return { type: 'class', value: 'bi bi-box-seam' };
+    }
+    return icon.type === 'img' && icon.value === 'llm_node.png'
+      ? { type: 'class', value: 'bi bi-box-seam' }
+      : icon;
+  }
+
   get inputs() {
     return Object.entries(this.data?.inputs ?? {})
       .filter(([key]) => key !== FLOW_DEPENDENCY_PORT_KEY)
diff --git a/src/app/shared/nodes/generic-node/generic-node.html b/src/app/shared/nodes/generic-node/generic-node.html
index 484c13c..4b977e9 100644
--- a/src/app/shared/nodes/generic-node/generic-node.html
+++ b/src/app/shared/nodes/generic-node/generic-node.html
@@ -23,10 +23,10 @@
   }
   <div class="llm-header" [class.cursor-move]="!isReadonly">
     <div class="llm-icon">
-      @if (isHumanNode()) {
-      <i class="bi bi-person-check-fill"></i>
+      @if (nodeIcon().type === 'class') {
+      <i [class]="nodeIcon().value"></i>
       } @else {
-      <img src="llm_node.png" alt="LLM" class="llm-icon-image" />
+      <img [src]="nodeIcon().value" alt="Node" class="llm-icon-image" />
       }
     </div>
     <div class="llm-header-content">
diff --git a/src/app/shared/nodes/generic-node/generic-node.ts b/src/app/shared/nodes/generic-node/generic-node.ts
index 37942e0..5da44bd 100644
--- a/src/app/shared/nodes/generic-node/generic-node.ts
+++ b/src/app/shared/nodes/generic-node/generic-node.ts
@@ -20,13 +20,19 @@ import {
   type UiConditionRule,
   evaluateUiConditionRule,
   flattenPrimitiveValues,
+  formatNodeTitle,
+  getOutputPillClass,
+  getOutputsTitle,
   getValueByPath,
+  isConditionalByPorts,
+  isHumanInteractiveNode,
   parentPath,
   pathToLabel,
   readUiConditionRule,
   readEffectiveUiVisibleConditionRule,
   readUiLabel,
   readUiGroup,
+  resolveNodeIcon,
   resolveSchemaRef,
   resolveSchemaPath,
   schemaFieldDescription,
@@ -436,34 +442,27 @@ export class GenericNodeComponent {
   }
 
   isHumanNode(): boolean {
-    return !!this.blockDescriptor?.interactionContract;
+    return isHumanInteractiveNode(this.blockDescriptor?.interactionContract);
   }
 
   isConditionalNode(): boolean {
-    const outputNames = this.resolvePorts('output').map((port) => port.name.trim().toLowerCase());
-    return outputNames.includes('true') && outputNames.includes('false');
+    return isConditionalByPorts(this.resolvePorts('output'));
   }
 
   nodeTitle(): string {
-    const type = this.blockType;
-    if (!type) return 'Node';
-    return type
-      .replace(/Block$/, '')
-      .replace(/([a-z0-9])([A-Z])/g, '$1 $2')
-      .trim();
+    return formatNodeTitle(this.blockType);
   }
 
   outputsTitle(): string {
-    return this.isConditionalNode() ? 'On Condition' : 'Outputs';
+    return getOutputsTitle(this.isConditionalNode());
   }
 
   outputPillClass(outputKey: string): string | null {
-    if (!this.isConditionalNode()) return null;
+    return getOutputPillClass(outputKey, this.isConditionalNode(), this.blockDescriptor?.schema);
+  }
 
-    const normalized = outputKey.trim().toLowerCase();
-    if (normalized === 'true') return 'llm-pill-output-true';
-    if (normalized === 'false') return 'llm-pill-output-false';
-    return null;
+  nodeIcon(): { type: 'class' | 'img'; value: string } {
+    return resolveNodeIcon(this.blockDescriptor?.schema, this.isHumanNode());
   }
 
   hasExecutionDependencyPorts(): boolean {
diff --git a/src/app/shared/nodes/node-utility.ts b/src/app/shared/nodes/node-utility.ts
index b127b1e..b4d1498 100644
--- a/src/app/shared/nodes/node-utility.ts
+++ b/src/app/shared/nodes/node-utility.ts
@@ -270,6 +270,59 @@ function isMeaningfullyPresent(value: unknown): boolean {
   return true;
 }
 
+export function isHumanInteractiveNode(interactionContract: unknown): boolean {
+  return !!interactionContract;
+}
+
+export function isConditionalByPorts(ports: { name: string }[]): boolean {
+  const names = ports.map((p) => p.name.trim().toLowerCase());
+  return names.includes('true') && names.includes('false');
+}
+
+export function formatNodeTitle(blockType: string | null | undefined, fallback = 'Node'): string {
+  if (!blockType) return fallback;
+  return blockType
+    .replace(/Block$/, '')
+    .replace(/([a-z0-9])([A-Z])/g, '$1 $2')
+    .trim();
+}
+
+export function getOutputsTitle(isConditional: boolean): string {
+  return isConditional ? 'On Condition' : 'Outputs';
+}
+
+export function getOutputPillClass(
+  outputKey: string,
+  isConditional: boolean,
+  schema?: Record<string, unknown> | null
+): string | null {
+  const outputStyles = schema?.['x-ui-output-styles'];
+  if (outputStyles && typeof outputStyles === 'object' && !Array.isArray(outputStyles)) {
+    const styleMap = outputStyles as Record<string, unknown>;
+    const style = styleMap[outputKey];
+    if (typeof style === 'string' && style.trim().length > 0) return style.trim();
+  }
+
+  if (!isConditional) return null;
+  const normalized = outputKey.trim().toLowerCase();
+  if (normalized === 'true') return 'llm-pill-output-true';
+  if (normalized === 'false') return 'llm-pill-output-false';
+  return null;
+}
+
+export function resolveNodeIcon(schema: Record<string, unknown> | null | undefined, hasInteractionContract: boolean): { type: 'class'; value: string } | { type: 'img'; value: string } {
+  const icon = schema?.['x-ui-icon'];
+  if (typeof icon === 'string' && icon.trim().length > 0) {
+    const trimmed = icon.trim();
+    if (trimmed.endsWith('.png') || trimmed.endsWith('.svg') || trimmed.endsWith('.jpg') || trimmed.endsWith('.webp')) {
+      return { type: 'img', value: trimmed };
+    }
+    return { type: 'class', value: trimmed };
+  }
+  if (hasInteractionContract) return { type: 'class', value: 'bi bi-person-check-fill' };
+  return { type: 'img', value: 'llm_node.png' };
+}
+
 export function splitTemplatedTextParts(text: string | null): TemplatedTextPart[] {
   if (!text) return [];
 
diff --git a/src/app/shared/nodes/task-step-node/task-step-node.html b/src/app/shared/nodes/task-step-node/task-step-node.html
index 7e54f34..5ef1319 100644
--- a/src/app/shared/nodes/task-step-node/task-step-node.html
+++ b/src/app/shared/nodes/task-step-node/task-step-node.html
@@ -47,10 +47,10 @@
 
   <div class="llm-header">
     <div class="llm-icon">
-      @if (isHumanNode()) {
-      <i class="bi bi-person-check-fill"></i>
+      @if (nodeIcon().type === 'class') {
+      <i [class]="nodeIcon().value"></i>
       } @else {
-      <img src="llm_node.png" alt="LLM" class="llm-icon-image" />
+      <img [src]="nodeIcon().value" alt="Node" class="llm-icon-image" />
       }
     </div>
     <div class="llm-header-content">
diff --git a/src/app/shared/nodes/task-step-node/task-step-node.ts b/src/app/shared/nodes/task-step-node/task-step-node.ts
index 4b030b7..4986927 100644
--- a/src/app/shared/nodes/task-step-node/task-step-node.ts
+++ b/src/app/shared/nodes/task-step-node/task-step-node.ts
@@ -13,12 +13,18 @@ import {
   type UiConditionRule,
   evaluateUiConditionRule,
   flattenPrimitiveValues,
+  formatNodeTitle,
+  getOutputPillClass,
+  getOutputsTitle,
+  isConditionalByPorts,
+  isHumanInteractiveNode,
   parentPath,
   pathToLabel,
   readUiConditionRule,
   readEffectiveUiVisibleConditionRule,
   readUiGroup,
   readUiLabel,
+  resolveNodeIcon,
   resolveSchemaRef,
   resolveSchemaPath,
   schemaFieldLabel,
@@ -242,34 +248,27 @@ export class TaskStepNodeComponent {
   }
 
   isHumanNode(): boolean {
-    return !!this.interactionContract();
+    return isHumanInteractiveNode(this.interactionContract());
   }
 
   isConditionalNode(): boolean {
-    const outputNames = this.resolvePorts('output').map((port) => port.name.trim().toLowerCase());
-    return outputNames.includes('true') && outputNames.includes('false');
+    return isConditionalByPorts(this.resolvePorts('output'));
   }
 
   nodeTitle(): string {
-    const type = this.blockType;
-    if (!type) return 'Task Step';
-    return type
-      .replace(/Block$/, '')
-      .replace(/([a-z0-9])([A-Z])/g, '$1 $2')
-      .trim();
+    return formatNodeTitle(this.blockType, 'Task Step');
   }
 
   outputsTitle(): string {
-    return this.isConditionalNode() ? 'On Condition' : 'Outputs';
+    return getOutputsTitle(this.isConditionalNode());
   }
 
   outputPillClass(outputKey: string): string | null {
-    if (!this.isConditionalNode()) return null;
+    return getOutputPillClass(outputKey, this.isConditionalNode(), this.blockDescriptor?.schema);
+  }
 
-    const normalized = outputKey.trim().toLowerCase();
-    if (normalized === 'true') return 'llm-pill-output-true';
-    if (normalized === 'false') return 'llm-pill-output-false';
-    return null;
+  nodeIcon(): { type: 'class' | 'img'; value: string } {
+    return resolveNodeIcon(this.blockDescriptor?.schema, this.isHumanNode());
   }
 
   inputDisplayLabel(inputKey: string): string {
@@ -566,30 +565,38 @@ export class TaskStepNodeComponent {
     return [{ text: value, isDynamicInput: false }];
   }
 
+  private schemaLoading = false;
+
   private async loadSchemaContext() {
+    if (this.schemaLoading) return;
     const type = this.blockType;
     if (!type) return;
 
-    const nodeFamily = this.resolvedNodeFamily();
-    const cachedDescriptor = nodeFamily === 'container'
-      ? this.containersService.peekContainerType(type)
-      : this.blocksService.peekBlockType(type);
-    const typeDescriptor = cachedDescriptor ?? (
-      nodeFamily === 'container'
-        ? await this.containersService.getContainerType(type)
-        : await this.blocksService.getBlockType(type)
-    );
-    this.blockDescriptor = (typeDescriptor ?? null) as BlockType | null;
-    this.blockSchema = (typeDescriptor?.schema ?? null) as Record<string, any> | null;
-    const typeKey = this.nodeTypeCacheKey();
-    if (typeKey) {
-      TaskStepNodeComponent.globalFieldSchemaCache.delete(typeKey);
-      TaskStepNodeComponent.globalFieldUiMetaCache.delete(typeKey);
-      TaskStepNodeComponent.globalFieldLabelCache.delete(typeKey);
+    this.schemaLoading = true;
+    try {
+      const nodeFamily = this.resolvedNodeFamily();
+      const cachedDescriptor = nodeFamily === 'container'
+        ? this.containersService.peekContainerType(type)
+        : this.blocksService.peekBlockType(type);
+      const typeDescriptor = cachedDescriptor ?? (
+        nodeFamily === 'container'
+          ? await this.containersService.getContainerType(type)
+          : await this.blocksService.getBlockType(type)
+      );
+      this.blockDescriptor = (typeDescriptor ?? null) as BlockType | null;
+      this.blockSchema = (typeDescriptor?.schema ?? null) as Record<string, any> | null;
+      const typeKey = this.nodeTypeCacheKey();
+      if (typeKey) {
+        TaskStepNodeComponent.globalFieldSchemaCache.delete(typeKey);
+        TaskStepNodeComponent.globalFieldUiMetaCache.delete(typeKey);
+        TaskStepNodeComponent.globalFieldLabelCache.delete(typeKey);
+      }
+      this.arrayFieldDefinitions = this.extractArrayFieldDefinitions(this.blockSchema);
+      this.rebuildDisplayState();
+      this.schemaReady = true;
+    } finally {
+      this.schemaLoading = false;
     }
-    this.arrayFieldDefinitions = this.extractArrayFieldDefinitions(this.blockSchema);
-    this.rebuildDisplayState();
-    this.schemaReady = true;
   }
 
   private interactionContract(): BlockInteractionContract | null {