diff --git a/packages/editor/src/editor.spec.ts b/packages/editor/src/editor.spec.ts
index 7fe986c0d3..647bd34a5f 100644
--- a/packages/editor/src/editor.spec.ts
+++ b/packages/editor/src/editor.spec.ts
@@ -518,6 +518,50 @@ describe('Editor', () => {
         const shellResult = editor._input.read().toString();
         expect(shellResult).to.be.equal(shellModifiedInput);
       });
+
+      it('returns a proper statement when editing previous code - input is not a statement', async() => {
+        const shellOriginalInput = 'foo';
+        const editorOutput = '20';
+        const shellModifiedInput = 'foo = 20';
+        const cmd = await fakeExternalEditor({
+          base: base.path,
+          name: 'editor-script.js',
+          output: editorOutput
+        });
+
+        editor = makeEditor({ cmd });
+
+        await editor.runEditCommand(shellOriginalInput);
+        let shellResult = editor._input.read().toString();
+        expect(shellResult).to.be.equal(shellModifiedInput);
+
+        await editor.runEditCommand('');
+        shellResult = editor._input.read().toString();
+        expect(shellResult).to.be.equal(shellModifiedInput);
+      });
+
+      it('returns a proper statement when editing previous code - input is a statement', async() => {
+        const shellOriginalInput = 'function () {}';
+        const editorOutput = `function () {
+          console.log(111);
+        }`;
+        const shellModifiedInput = 'function () { console.log(111); }';
+        const cmd = await fakeExternalEditor({
+          base: base.path,
+          name: 'editor-script.js',
+          output: editorOutput
+        });
+
+        editor = makeEditor({ cmd });
+
+        await editor.runEditCommand(shellOriginalInput);
+        let shellResult = editor._input.read().toString();
+        expect(shellResult).to.be.equal(shellModifiedInput);
+
+        await editor.runEditCommand('');
+        shellResult = editor._input.read().toString();
+        expect(shellResult).to.be.equal(shellModifiedInput);
+      });
     });
   });
 });
diff --git a/packages/editor/src/editor.ts b/packages/editor/src/editor.ts
index c36a7eeb37..b3ec5d873e 100644
--- a/packages/editor/src/editor.ts
+++ b/packages/editor/src/editor.ts
@@ -28,6 +28,7 @@ export class Editor {
   _instanceState: ShellInstanceState;
   _loadExternalCode: (input: string, filename: string) => Promise<ShellResult>;
   _lastContent: string;
+  _lastInputCode: string;
   print: (...args: any[]) => Promise<void>;
 
   constructor({ input, vscodeDir, tmpDir, instanceState, loadExternalCode }: EditorOptions) {
@@ -37,6 +38,7 @@ export class Editor {
     this._instanceState = instanceState;
     this._loadExternalCode = loadExternalCode;
     this._lastContent = '';
+    this._lastInputCode = '';
     this.print = instanceState.context.print;
 
     // Add edit command support to shell api.
@@ -173,6 +175,19 @@ export class Editor {
     return `${originalCode} = ${modifiedCode}`;
   }
 
+  _setLastInputCode(code: string): void {
+    if (code !== '') {
+      this._lastInputCode = code;
+    }
+  }
+
+  _getLastInputCode(code: string): string {
+    if (code !== '') {
+      return code;
+    }
+    return this._lastInputCode;
+  }
+
   async runEditCommand(code: string): Promise<void> {
     await this.print('Opening an editor...');
 
@@ -183,6 +198,8 @@ export class Editor {
       throw new Error('Command failed with an error: please define an external editor');
     }
 
+    this._setLastInputCode(code);
+
     const content = await this._getEditorContent(code);
     const ext = await this._getExtension(editor);
     const tmpDoc = await this._createTempFile({ content, ext });
@@ -208,7 +225,7 @@ export class Editor {
 
       if (exitCode === 0) {
         const modifiedCode = await this._readAndDeleteTempFile(tmpDoc);
-        const result = this._prepareResult({ originalCode: code, modifiedCode });
+        const result = this._prepareResult({ originalCode: this._getLastInputCode(code), modifiedCode });
 
         // Write a content from the editor to the parent readable stream.
         this._input.unshift(result);