From a68606dc6bfed4c7f839bf2ea1252c834a23c4fd Mon Sep 17 00:00:00 2001
From: Sean Doyle <sean.p.doyle24@gmail.com>
Date: Tue, 16 Aug 2022 11:29:15 -0400
Subject: [PATCH] Turbo Streams: Manage element focus

When a `<turbo-stream>` modifies the document, it has the potential to
affect which element has focus.

For example, consider an element with an `[id]` that has focus:

```html
<label for="an-input">
<input id="an-input" value="an invalid value"> <!-- this element has focus -->
```

Next, consider a `<turbo-stream>` element to replace it:

```html
<turbo-stream action="replace" target="an-input">
  <template>
    <input id="an-input" value="an invalid value" class="invalid-input">
  </template>
</turbo-stream>
```

Prior to this commit, rendering that `<turbo-stream>` would remove the
element with focus, and never restore it.

After this commit, the `Session` will capture the `[id]` value of the
element with focus (if there is any), then "restore" focus to an element
in the document with a matching `[id]` attribute _after_ the render.

Similarly, consider a `<turbo-stream>` that appends an element with
`[autofocus]`:

```html
<turbo-stream action="append" targets="body">
  <template>
    <input autofocus>
  </template>
</turbo-stream>
```

Prior to this commit, inserting an `[autofocus]` into the document with
a `<turbo-stream>` had no effect.

After this commit, the `Session` will scan any `<turbo-stream>` elements
its about to render, extracting the first focusable element that
declares an `[autofocus]` attribute.

Once the rendering is complete, it will attempt to autofocus that
element. Several scenarios will prevent that, including:

* there aren't any `[autofocus]` elements in the collection of
  `<turbo-stream>` elements
* the `[autofocus]` element does not exist in the document after the
  rendering is complete
* the document already has an element with focus
---
 src/core/renderer.ts                        |  6 +-
 src/core/snapshot.ts                        | 11 +---
 src/core/streams/stream_message_renderer.ts | 63 ++++++++++++++++++++-
 src/tests/fixtures/stream.html              |  4 ++
 src/tests/functional/autofocus_tests.ts     | 58 +++++++++++++++++++
 src/tests/functional/stream_tests.ts        | 44 +++++++++++++-
 src/util.ts                                 | 31 ++++++++++
 7 files changed, 199 insertions(+), 18 deletions(-)

diff --git a/src/core/renderer.ts b/src/core/renderer.ts
index e77be9d34..ce64b2bbe 100644
--- a/src/core/renderer.ts
+++ b/src/core/renderer.ts
@@ -55,7 +55,7 @@ export abstract class Renderer<E extends Element, S extends Snapshot<E> = Snapsh
 
   focusFirstAutofocusableElement() {
     const element = this.connectedSnapshot.firstAutofocusableElement
-    if (elementIsFocusable(element)) {
+    if (element) {
       element.focus()
     }
   }
@@ -94,7 +94,3 @@ export abstract class Renderer<E extends Element, S extends Snapshot<E> = Snapsh
     return this.currentSnapshot.getPermanentElementMapForSnapshot(this.newSnapshot)
   }
 }
-
-function elementIsFocusable(element: any): element is { focus: () => void } {
-  return element && typeof element.focus == "function"
-}
diff --git a/src/core/snapshot.ts b/src/core/snapshot.ts
index 7d5431117..8bba5d206 100644
--- a/src/core/snapshot.ts
+++ b/src/core/snapshot.ts
@@ -1,3 +1,5 @@
+import { queryAutofocusableElement } from "../util"
+
 export class Snapshot<E extends Element = Element> {
   readonly element: E
 
@@ -26,14 +28,7 @@ export class Snapshot<E extends Element = Element> {
   }
 
   get firstAutofocusableElement() {
-    const inertDisabledOrHidden = "[inert], :disabled, [hidden], details:not([open]), dialog:not([open])"
-
-    for (const element of this.element.querySelectorAll("[autofocus]")) {
-      if (element.closest(inertDisabledOrHidden) == null) return element
-      else continue
-    }
-
-    return null
+    return queryAutofocusableElement(this.element)
   }
 
   get permanentElements() {
diff --git a/src/core/streams/stream_message_renderer.ts b/src/core/streams/stream_message_renderer.ts
index 549e4de8d..e7158b74a 100644
--- a/src/core/streams/stream_message_renderer.ts
+++ b/src/core/streams/stream_message_renderer.ts
@@ -2,12 +2,17 @@ import { StreamMessage } from "./stream_message"
 import { StreamElement } from "../../elements/stream_element"
 import { Bardo, BardoDelegate } from "../bardo"
 import { PermanentElementMap, getPermanentElementById, queryPermanentElementsAll } from "../snapshot"
+import { around, elementIsFocusable, queryAutofocusableElement, waitForCallback } from "../../util"
 
 export class StreamMessageRenderer implements BardoDelegate {
   render({ fragment }: StreamMessage) {
-    Bardo.preservingPermanentElements(this, getPermanentElementMapForFragment(fragment), () =>
-      document.documentElement.appendChild(fragment)
-    )
+    Bardo.preservingPermanentElements(this, getPermanentElementMapForFragment(fragment), () => {
+      withAutofocusFromFragment(fragment, () => {
+        withPreservedFocus(() => {
+          document.documentElement.appendChild(fragment)
+        })
+      })
+    })
   }
 
   enteringBardo(currentPermanentElement: Element, newPermanentElement: Element) {
@@ -34,3 +39,55 @@ function getPermanentElementMapForFragment(fragment: DocumentFragment): Permanen
 
   return permanentElementMap
 }
+
+async function withAutofocusFromFragment(fragment: DocumentFragment, callback: () => void) {
+  const generatedID = `turbo-stream-autofocus-${Date.now()}`
+  const turboStreams = fragment.querySelectorAll<StreamElement>("turbo-stream")
+  const elementWithAutofocus = firstAutofocusableElementInStreams(turboStreams)
+  let willAutofocus = generatedID
+
+  if (elementWithAutofocus) {
+    if (elementWithAutofocus.id) willAutofocus = elementWithAutofocus.id
+
+    elementWithAutofocus.id = willAutofocus
+  }
+
+  await waitForCallback(callback)
+
+  const hasNoActiveElement = document.activeElement == null || document.activeElement == document.body
+
+  if (hasNoActiveElement && willAutofocus) {
+    const elementToAutofocus = document.getElementById(willAutofocus)
+
+    if (elementIsFocusable(elementToAutofocus)) {
+      elementToAutofocus.focus()
+    }
+    if (elementToAutofocus && elementToAutofocus.id == generatedID) {
+      elementToAutofocus.removeAttribute("id")
+    }
+  }
+}
+
+async function withPreservedFocus(callback: () => void) {
+  const [activeElementBeforeRender, activeElementAfterRender] = await around(callback, () => document.activeElement)
+
+  const restoreFocusTo = activeElementBeforeRender && activeElementBeforeRender.id
+
+  if (restoreFocusTo) {
+    const elementToFocus = document.getElementById(restoreFocusTo)
+
+    if (elementIsFocusable(elementToFocus) && elementToFocus != activeElementAfterRender) {
+      elementToFocus.focus()
+    }
+  }
+}
+
+function firstAutofocusableElementInStreams(nodes: NodeListOf<StreamElement>) {
+  for (const stream of nodes) {
+    const elementWithAutofocus = queryAutofocusableElement(stream.templateElement.content)
+
+    if (elementWithAutofocus) return elementWithAutofocus
+  }
+
+  return null
+}
diff --git a/src/tests/fixtures/stream.html b/src/tests/fixtures/stream.html
index 164f45256..85e8c94e4 100644
--- a/src/tests/fixtures/stream.html
+++ b/src/tests/fixtures/stream.html
@@ -35,5 +35,9 @@
     <div id="messages_3" class="messages">
       <div class="message">Third</div>
     </div>
+
+    <div id="container">
+      <input id="container-element">
+    </div>
   </body>
 </html>
diff --git a/src/tests/functional/autofocus_tests.ts b/src/tests/functional/autofocus_tests.ts
index 06fa68e46..e69d5e5a1 100644
--- a/src/tests/functional/autofocus_tests.ts
+++ b/src/tests/functional/autofocus_tests.ts
@@ -108,3 +108,61 @@ test("test navigating a frame with a turbo-frame targeting the frame autofocuses
     "focuses the first [autofocus] element in frame"
   )
 })
+
+test("test receiving a Turbo Stream message with an [autofocus] element when the activeElement is the document", async ({
+  page,
+}) => {
+  await page.evaluate(() => {
+    if (document.activeElement instanceof HTMLElement) {
+      document.activeElement.blur()
+    }
+    window.Turbo.renderStreamMessage(`
+      <turbo-stream action="append" targets="body">
+        <template><input id="autofocus-from-stream" autofocus></template>
+      </turbo-stream>
+    `)
+  })
+  await nextBeat()
+
+  assert.ok(
+    await hasSelector(page, "#autofocus-from-stream:focus"),
+    "focuses the [autofocus] element in from the turbo-stream"
+  )
+})
+
+test("test autofocus from a Turbo Stream message does not leak a placeholder [id]", async ({ page }) => {
+  await page.evaluate(() => {
+    if (document.activeElement instanceof HTMLElement) {
+      document.activeElement.blur()
+    }
+    window.Turbo.renderStreamMessage(`
+      <turbo-stream action="append" targets="body">
+        <template><div id="container-from-stream"><input autofocus></div></template>
+      </turbo-stream>
+    `)
+  })
+  await nextBeat()
+
+  assert.ok(
+    await hasSelector(page, "#container-from-stream input:focus"),
+    "focuses the [autofocus] element in from the turbo-stream"
+  )
+})
+
+test("test receiving a Turbo Stream message with an [autofocus] element when an element within the document has focus", async ({
+  page,
+}) => {
+  await page.evaluate(() => {
+    window.Turbo.renderStreamMessage(`
+      <turbo-stream action="append" targets="body">
+        <template><input id="autofocus-from-stream" autofocus></template>
+      </turbo-stream>
+    `)
+  })
+  await nextBeat()
+
+  assert.ok(
+    await hasSelector(page, "#first-autofocus-element:focus"),
+    "focuses the first [autofocus] element on the page"
+  )
+})
diff --git a/src/tests/functional/stream_tests.ts b/src/tests/functional/stream_tests.ts
index 356ece95f..5a54a57a9 100644
--- a/src/tests/functional/stream_tests.ts
+++ b/src/tests/functional/stream_tests.ts
@@ -1,6 +1,6 @@
 import { test } from "@playwright/test"
 import { assert } from "chai"
-import { nextBeat, nextEventNamed, readEventLogs } from "../helpers/page"
+import { hasSelector, nextBeat, nextEventNamed, readEventLogs } from "../helpers/page"
 
 test.beforeEach(async ({ page }) => {
   await page.goto("/src/tests/fixtures/stream.html")
@@ -50,7 +50,7 @@ test("test receiving a message without a template", async ({ page }) => {
     `)
   )
 
-  assert.equal(await await page.locator("#messages").count(), 0, "removes target element")
+  assert.equal(await page.locator("#messages").count(), 0, "removes target element")
 })
 
 test("test receiving a message with a <script> element", async ({ page }) => {
@@ -123,3 +123,43 @@ test("test receiving a stream message asynchronously", async ({ page }) => {
 
   assert.deepEqual(await messages.allTextContents(), ["First", "Hello world!"])
 })
+
+test("test receiving an update stream message preserves focus if the activeElement has an [id]", async ({ page }) => {
+  await page.locator("input#container-element").focus()
+  await page.evaluate(() => {
+    window.Turbo.renderStreamMessage(`
+      <turbo-stream action="update" target="container">
+        <template><textarea id="container-element"></textarea></template>
+      </turbo-stream>
+    `)
+  })
+  await nextBeat()
+
+  assert.ok(await hasSelector(page, "textarea#container-element:focus"))
+})
+
+test("test receiving a replace stream message preserves focus if the activeElement has an [id]", async ({ page }) => {
+  await page.locator("input#container-element").focus()
+  await page.evaluate(() => {
+    window.Turbo.renderStreamMessage(`
+      <turbo-stream action="replace" target="container-element">
+        <template><textarea id="container-element"></textarea></template>
+      </turbo-stream>
+    `)
+  })
+  await nextBeat()
+
+  assert.ok(await hasSelector(page, "textarea#container-element:focus"))
+})
+
+test("test receiving a remove stream message preserves focus blurs the activeElement", async ({ page }) => {
+  await page.locator("#container-element").focus()
+  await page.evaluate(() => {
+    window.Turbo.renderStreamMessage(`
+      <turbo-stream action="remove" target="container-element"></turbo-stream>
+    `)
+  })
+  await nextBeat()
+
+  assert.notOk(await hasSelector(page, ":focus"))
+})
diff --git a/src/util.ts b/src/util.ts
index dbd5c1aa6..1509660ce 100644
--- a/src/util.ts
+++ b/src/util.ts
@@ -188,3 +188,34 @@ export function setMetaContent(name: string, content: string) {
 
   return element
 }
+
+type FocusableElement = Element & { focus: () => void }
+
+export function elementIsFocusable(element: Element | null): element is FocusableElement {
+  const inertDisabledOrHidden = "[inert], :disabled, [hidden], details:not([open]), dialog:not([open])"
+
+  return !!element && element.closest(inertDisabledOrHidden) == null && typeof (element as any).focus == "function"
+}
+
+export function queryAutofocusableElement(root: Element | DocumentFragment): FocusableElement | null {
+  for (const element of root.querySelectorAll("[autofocus]")) {
+    if (elementIsFocusable(element)) return element
+    else continue
+  }
+
+  return null
+}
+
+export async function around<T>(callback: () => void, reader: () => T): Promise<[T, T]> {
+  const before = reader()
+  await waitForCallback(callback)
+  const after = reader()
+
+  return [before, after]
+}
+
+export function waitForCallback(callback: () => void): Promise<void> {
+  callback()
+
+  return nextAnimationFrame()
+}