3 files changed, 9 insertions, 4 deletions
diff --git a/src/client/views/ScriptBox.tsx b/src/client/views/ScriptBox.tsx
index 8f08224c8..8f06cf770 100644
--- a/src/client/views/ScriptBox.tsx
+++ b/src/client/views/ScriptBox.tsx
@@ -99,6 +99,9 @@ export class ScriptBox extends React.Component<ScriptBoxProps> {
         let params: string[] = [];
         let setParams = (p: string[]) => params.splice(0, params.length, ...p);
         let scriptingBox = <ScriptBox initialText={originalText} setParams={setParams} onCancel={() => overlayDisposer()} onSave={(text, onError) => {
+            if (prewrapper) {
+                text = prewrapper + text + (postwrapper ? postwrapper : "");
+            }
             const script = CompileScript(text, {
                 params: { this: Doc.name },
                 typecheck: false,
diff --git a/src/client/views/collections/CollectionBaseView.tsx b/src/client/views/collections/CollectionBaseView.tsx
index bd8d56851..5829f0626 100644
--- a/src/client/views/collections/CollectionBaseView.tsx
+++ b/src/client/views/collections/CollectionBaseView.tsx
@@ -104,9 +104,6 @@ export class CollectionBaseView extends React.Component<CollectionViewProps> {
         if (this.props.fieldExt) { // bcz: fieldExt !== undefined means this is an overlay layer
             Doc.GetProto(doc).annotationOn = this.props.Document;
         }
-        if (doc.type === DocumentType.BUTTON) {
-            doc.collectionContext = this.props.Document;  // used by docList() function in Doc.ts so that buttons can iterate over the documents in their collection
-        }
         allowDuplicates = true;
         let targetDataDoc = this.props.fieldExt || this.props.Document.isTemplate ? this.extensionDoc : this.props.Document;
         let targetField = (this.props.fieldExt || this.props.Document.isTemplate) && this.props.fieldExt ? this.props.fieldExt : this.props.fieldKey;
diff --git a/src/client/views/nodes/DocumentView.tsx b/src/client/views/nodes/DocumentView.tsx
index 9c2cf5f01..940a66b36 100644
--- a/src/client/views/nodes/DocumentView.tsx
+++ b/src/client/views/nodes/DocumentView.tsx
@@ -643,7 +643,12 @@ export class DocumentView extends DocComponent<DocumentViewProps, Document>(Docu
         let onClicks: ContextMenuProps[] = existingOnClick && "subitems" in existingOnClick ? existingOnClick.subitems : [];
         onClicks.push({ description: this.props.Document.isButton || this.props.Document.onClick ? "Remove Click Behavior" : "Follow Link", event: this.makeBtnClicked, icon: "concierge-bell" });
         onClicks.push({ description: "Edit onClick Script", icon: "edit", event: (obj: any) => ScriptBox.EditButtonScript("On Button Clicked ...", this.props.Document, "onClick", obj.x, obj.y) });
-        onClicks.push({ description: "Edit onClick Foreach Doc Script", icon: "edit", event: (obj: any) => ScriptBox.EditButtonScript("Foreach Collection Doc (d) => ", this.props.Document, "onClick", obj.x, obj.y, "docList(this.collectionContext.data).map(d => {", "});\n") });
+        onClicks.push({
+            description: "Edit onClick Foreach Doc Script", icon: "edit", event: (obj: any) => {
+                this.props.Document.collectionContext = this.props.ContainingCollectionView && this.props.ContainingCollectionView.props.Document;
+                ScriptBox.EditButtonScript("Foreach Collection Doc (d) => ", this.props.Document, "onClick", obj.x, obj.y, "docList(this.collectionContext.data).map(d => {", "});\n");
+            }
+        });
         !existingOnClick && cm.addItem({ description: "OnClick...", subitems: onClicks, icon: "hand-point-right" });
 
         let existing = ContextMenu.Instance.findByDescription("Layout...");