1 files changed, 4 insertions, 2 deletions
diff --git a/src/client/views/nodes/DiagramBox.tsx b/src/client/views/nodes/DiagramBox.tsx
index 7cfccf0dc..6a31f64ce 100644
--- a/src/client/views/nodes/DiagramBox.tsx
+++ b/src/client/views/nodes/DiagramBox.tsx
@@ -185,6 +185,8 @@ export class DiagramBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         return '( )';
     };
 
+    setRef = (r: HTMLDivElement | null) => this.fixWheelEvents(r, this._props.isContentActive);
+    setDiagramBoxRef = (r: HTMLDivElement | null) => r && this.renderMermaidAsync.call(this, this.removeWords(this.mermaidcode), r);
     render() {
         return (
             <div
@@ -192,7 +194,7 @@ export class DiagramBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 style={{
                     pointerEvents: this._props.isContentActive() ? undefined : 'none',
                 }}
-                ref={r => this.fixWheelEvents(r, this._props.isContentActive)}>
+                ref={this.setRef}>
                 <div className="DIYNodeBox-searchbar">
                     <input type="text" value={this._inputValue} onKeyDown={action(e => e.key === 'Enter' && this.generateMermaidCode())} onChange={action(e => (this._inputValue = e.target.value))} />
                     <button type="button" onClick={this.generateMermaidCode}>
@@ -208,7 +210,7 @@ export class DiagramBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     ) : this._generating ? (
                         <div className="loading-circle" />
                     ) : (
-                        <div className="diagramBox" ref={r => r && this.renderMermaidAsync.call(this, this.removeWords(this.mermaidcode), r)}>
+                        <div className="diagramBox" ref={this.setDiagramBoxRef}>
                             {this._errorMessage || 'Type a prompt to generate a diagram'}
                         </div>
                     )}