Fix bug with missing draw shape after undo

packages/tldraw/src/state/tools/DrawTool/DrawTool.spec.ts

@@ -1,4 +1,6 @@
 import { TldrawApp } from '~state'
+import { TldrawTestApp } from '~test'
+import { TDShapeType } from '~types'
 import { DrawTool } from '.'
 
 describe('DrawTool', () => {
@@ -7,3 +9,61 @@ describe('DrawTool', () => {
     new DrawTool(app)
   })
 })
+
+describe('When shift+clicking to extend a shape', () => {
+  it('extends the same shape', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas([0, 0])
+    app.movePointer([100, 100])
+    app.movePointer([200, 200])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+  })
+
+  it('does not extend after switching tools the same shape', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas([0, 0])
+    app.movePointer([100, 100])
+    app.movePointer([200, 200])
+    app.stopPointing()
+    app.selectTool('select')
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(2)
+  })
+
+  it('does not extend after undo', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas([0, 0])
+    app.movePointer([100, 100])
+    app.movePointer([200, 200])
+    app.stopPointing()
+    app.undo()
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+  })
+
+  it('does not extend if no shape is present', () => {
+    const app = new TldrawTestApp()
+    app.reset()
+    app.selectTool(TDShapeType.Draw)
+    app.pointCanvas({ x: 300, y: 300, shiftKey: true })
+    app.movePointer([400, 400])
+    app.stopPointing()
+    expect(app.shapes.length).toBe(1)
+  })
+})

packages/tldraw/src/state/tools/DrawTool/DrawTool.ts

@@ -42,9 +42,10 @@ export class DrawTool extends BaseTool {
       currentPoint,
       appState: { currentPageId, currentStyle },
     } = this.app
-    if (info.shiftKey && this.lastShapeId) {
+    const previous = this.lastShapeId && this.app.getShape(this.lastShapeId)
+    if (info.shiftKey && previous) {
       // Extend the previous shape
-      this.app.startSession(SessionType.Draw, this.lastShapeId)
+      this.app.startSession(SessionType.Draw, previous.id)
       this.setStatus(Status.Extending)
     } else {
       // Create a new shape