diff --git a/ui/src/flux/components/ExpressionNode.tsx b/ui/src/flux/components/ExpressionNode.tsx
index 6bb1dda1e..5fe996ea3 100644
--- a/ui/src/flux/components/ExpressionNode.tsx
+++ b/ui/src/flux/components/ExpressionNode.tsx
@@ -116,7 +116,8 @@ class ExpressionNode extends PureComponent<Props, State> {
                     onDelete={onDeleteFuncNode}
                     onToggleYield={onToggleYield}
                     isYieldable={isAfterFilter && isAfterRange}
-                    isYielding={this.isBeforeFuncYield(i)}
+                    isYielding={this.isBeforeYielding(i)}
+                    isYieldedInScript={this.isYieldNodeIndex(i + 1)}
                     declarationID={declarationID}
                     onGenerateScript={onGenerateScript}
                     declarationsFromBody={declarationsFromBody}
@@ -125,7 +126,10 @@ class ExpressionNode extends PureComponent<Props, State> {
                   />
                 )
 
-                if (nonYieldableIndexesToggled[i]) {
+                if (
+                  nonYieldableIndexesToggled[i] &&
+                  !this.isYieldNodeIndex(i + 1)
+                ) {
                   const script = scriptUpToYield(
                     bodyID,
                     declarationID,
@@ -133,10 +137,6 @@ class ExpressionNode extends PureComponent<Props, State> {
                     false
                   )
 
-                  if (this.isBeforeFuncYield(i)) {
-                    return funcNode
-                  }
-
                   return (
                     <Fragment key={`${i}-notInScript`}>
                       {funcNode}
@@ -167,7 +167,8 @@ class ExpressionNode extends PureComponent<Props, State> {
                         onDelete={onDeleteFuncNode}
                         onToggleYield={this.handleHideImplicitYield}
                         isYieldable={isAfterFilter && isAfterRange}
-                        isYielding={this.isBeforeFuncYield(i)}
+                        isYielding={this.isBeforeYielding(i)}
+                        isYieldedInScript={this.isYieldNodeIndex(i + 1)}
                         declarationID={declarationID}
                         onGenerateScript={onGenerateScript}
                         declarationsFromBody={declarationsFromBody}
@@ -202,9 +203,14 @@ class ExpressionNode extends PureComponent<Props, State> {
     )
   }
 
-  private isBeforeFuncYield(funcIndex: number): boolean {
+  private isBeforeYielding(funcIndex: number): boolean {
     const {funcs, isLastBody} = this.props
-    const {isImplicitYieldToggled} = this.state
+    const {isImplicitYieldToggled, nonYieldableIndexesToggled} = this.state
+    const beforeToggledLastYield = !!nonYieldableIndexesToggled[funcIndex]
+
+    if (beforeToggledLastYield) {
+      return true
+    }
 
     if (
       funcIndex === funcs.length - 1 &&
@@ -218,9 +224,14 @@ class ExpressionNode extends PureComponent<Props, State> {
       return false
     }
 
-    const nextFunc = funcs[funcIndex + 1]
+    return this.isYieldNodeIndex(funcIndex + 1)
+  }
 
-    if (nextFunc.name === 'yield') {
+  private isYieldNodeIndex(funcIndex) {
+    const {funcs} = this.props
+    const nextFunc = _.get(funcs, `${funcIndex}`, null)
+
+    if (nextFunc && nextFunc.name === 'yield') {
       return true
     }
 
diff --git a/ui/src/flux/components/FuncNode.tsx b/ui/src/flux/components/FuncNode.tsx
index 56c3f2458..e116dd249 100644
--- a/ui/src/flux/components/FuncNode.tsx
+++ b/ui/src/flux/components/FuncNode.tsx
@@ -30,6 +30,7 @@ interface Props {
   isYielding: boolean
   isYieldable: boolean
   onDeleteBody: (bodyID: string) => void
+  isYieldedInScript: boolean
 }
 
 interface State {
@@ -182,19 +183,14 @@ export default class FuncNode extends PureComponent<Props, State> {
       index,
       bodyID,
       declarationID,
-      isYieldable,
       onToggleYieldWithLast,
-      isYielding,
+      isYieldedInScript,
     } = this.props
 
-    if (isYieldable) {
+    if (isYieldedInScript) {
       onToggleYield(bodyID, declarationID, index)
     } else {
-      if (isYielding) {
-        onToggleYield(bodyID, declarationID, index)
-      } else {
-        onToggleYieldWithLast(index)
-      }
+      onToggleYieldWithLast(index)
     }
   }