google · BenHenning · Apr 23, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 21, 2025
@@ -60,6 +60,7 @@ export class FocusManager {
   registeredTrees: Array<IFocusableTree> = [];
 
   private currentlyHoldsEphemeralFocus: boolean = false;
+  private lockFocusStateChanges: boolean = false;
 
   constructor(
     addGlobalEventListener: (type: string, listener: EventListener) => void,
@@ -89,7 +90,16 @@ export class FocusManager {
       }
 
       if (newNode) {
-        this.focusNode(newNode);
+        const newTree = newNode.getFocusableTree();
+        const oldTree = this.focusedNode?.getFocusableTree();
+        if (newNode === newTree.getRootFocusableNode() && newTree !== oldTree) {
+          // If the root of the tree is the one taking focus (such as due to
+          // being tabbed), try to focus the whole tree explicitly to ensure the
+          // correct node re-receives focus.
+          this.focusTree(newTree);
+        } else {
+          this.focusNode(newNode);
+        }
       } else {
         this.defocusCurrentFocusedNode();
       }
@@ -108,6 +118,7 @@ export class FocusManager {
    * certain whether the tree has been registered.
    */
   registerTree(tree: IFocusableTree): void {
+    this.ensureManagerIsUnlocked();
     if (this.isRegistered(tree)) {
       throw Error(`Attempted to re-register already registered tree: ${tree}.`);
     }
@@ -133,10 +144,11 @@ export class FocusManager {
    * this manager.
    */
   unregisterTree(tree: IFocusableTree): void {
+    this.ensureManagerIsUnlocked();
     if (!this.isRegistered(tree)) {
       throw Error(`Attempted to unregister not registered tree: ${tree}.`);
     }
-    const treeIndex = this.registeredTrees.findIndex((tree) => tree === tree);
+    const treeIndex = this.registeredTrees.findIndex((reg) => reg === tree);
     this.registeredTrees.splice(treeIndex, 1);
 
     const focusedNode = FocusableTreeTraverser.findFocusedNode(tree);
@@ -192,11 +204,14 @@ export class FocusManager {
    *     focus.
    */
   focusTree(focusableTree: IFocusableTree): void {
+    this.ensureManagerIsUnlocked();
     if (!this.isRegistered(focusableTree)) {
       throw Error(`Attempted to focus unregistered tree: ${focusableTree}.`);
     }
     const currNode = FocusableTreeTraverser.findFocusedNode(focusableTree);
-    this.focusNode(currNode ?? focusableTree.getRootFocusableNode());
+    const nodeToRestore = focusableTree.getRestoredFocusableNode(currNode);
+    const rootFallback = focusableTree.getRootFocusableNode();
+    this.focusNode(nodeToRestore ?? currNode ?? rootFallback);
   }
 
   /**
@@ -205,18 +220,37 @@ export class FocusManager {
    * Any previously focused node will be updated to be passively highlighted (if
    * it's in a different focusable tree) or blurred (if it's in the same one).
    *
-   * @param focusableNode The node that should receive active
-   *     focus.
+   * @param focusableNode The node that should receive active focus.
    */
   focusNode(focusableNode: IFocusableNode): void {
+    this.ensureManagerIsUnlocked();
+    if (this.focusedNode === focusableNode) return; // State is unchanged.
+
     const nextTree = focusableNode.getFocusableTree();
     if (!this.isRegistered(nextTree)) {
       throw Error(`Attempted to focus unregistered node: ${focusableNode}.`);
     }
+
+    // Safety check for ensuring focusNode() doesn't get called for a node that
+    // isn't actually hooked up to its parent tree correctly (since this can
+    // cause weird inconsistencies).
+    const matchedNode = FocusableTreeTraverser.findFocusableNodeFor(
+      focusableNode.getFocusableElement(),
+      nextTree,
+    );
+    if (matchedNode !== focusableNode) {
+      throw Error(
+        `Attempting to focus node which isn't recognized by its parent tree: ` +
+          `${focusableNode}.`,
+      );
+    }
+
     const prevNode = this.focusedNode;
-    if (prevNode && prevNode.getFocusableTree() !== nextTree) {
-      this.setNodeToPassive(prevNode);
+    const prevTree = prevNode?.getFocusableTree();
+    if (prevNode && prevTree !== nextTree) {
+      this.passivelyFocusNode(prevNode, nextTree);
     }
+
     // If there's a focused node in the new node's tree, ensure it's reset.
     const prevNodeNextTree = FocusableTreeTraverser.findFocusedNode(nextTree);
     const nextTreeRoot = nextTree.getRootFocusableNode();
@@ -229,9 +263,10 @@ export class FocusManager {
     if (nextTreeRoot !== focusableNode) {
       this.removeHighlight(nextTreeRoot);
     }
+
     if (!this.currentlyHoldsEphemeralFocus) {
       // Only change the actively focused node if ephemeral state isn't held.
-      this.setNodeToActive(focusableNode);
+      this.activelyFocusNode(focusableNode, prevTree ?? null);
     }
     this.focusedNode = focusableNode;
   }
@@ -257,6 +292,7 @@ export class FocusManager {
   takeEphemeralFocus(
     focusableElement: HTMLElement | SVGElement,
   ): ReturnEphemeralFocus {
+    this.ensureManagerIsUnlocked();
     if (this.currentlyHoldsEphemeralFocus) {
       throw Error(
         `Attempted to take ephemeral focus when it's already held, ` +
@@ -266,7 +302,7 @@ export class FocusManager {
     this.currentlyHoldsEphemeralFocus = true;
 
     if (this.focusedNode) {
-      this.setNodeToPassive(this.focusedNode);
+      this.passivelyFocusNode(this.focusedNode, null);
     }
     focusableElement.focus();
 
@@ -282,34 +318,124 @@ export class FocusManager {
       this.currentlyHoldsEphemeralFocus = false;
 
       if (this.focusedNode) {
-        this.setNodeToActive(this.focusedNode);
+        this.activelyFocusNode(this.focusedNode, null);
       }
     };
   }
 
+  /**
+   * Ensures that the manager is currently allowing operations that change its
+   * internal focus state (such as via focusNode()).
+   *
+   * If the manager is currently not allowing state changes, an exception is
+   * thrown.
+   */
+  private ensureManagerIsUnlocked(): void {
+    if (this.lockFocusStateChanges) {
+      throw Error(
+        'FocusManager state changes cannot happen in a tree/node focus/blur ' +
+          'callback.',
+      );
+    }
+  }
+
+  /**
+   * Defocuses the current actively focused node tracked by the manager, iff
+   * there's a node being tracked and the manager doesn't have ephemeral focus.
+   */
   private defocusCurrentFocusedNode(): void {
     // The current node will likely be defocused while ephemeral focus is held,
     // but internal manager state shouldn't change since the node should be
     // restored upon exiting ephemeral focus mode.
     if (this.focusedNode && !this.currentlyHoldsEphemeralFocus) {
-      this.setNodeToPassive(this.focusedNode);
+      this.passivelyFocusNode(this.focusedNode, null);
       this.focusedNode = null;
     }
   }
 
-  private setNodeToActive(node: IFocusableNode): void {
+  /**
+   * Marks the specified node as actively focused, also calling related lifecycle
+   * callback methods for both the node and its parent tree. This ensures that
+   * the node is properly styled to indicate its active focus.
+   *
+   * This does not change the manager's currently tracked node, nor does it
+   * change any other nodes.
+   *
+   * @param node The node to be actively focused.
+   * @param prevTree The tree of the previously actively focused node, or null
+   *     if there wasn't a previously actively focused node.
+   */
+  private activelyFocusNode(
+    node: IFocusableNode,
+    prevTree: IFocusableTree | null,
+  ): void {
+    // Note that order matters here. Focus callbacks are allowed to change
+    // element visibility which can influence focusability, including for a
+    // node's focusable element (which *is* allowed to be invisible until the
+    // node needs to be focused).
+    this.lockFocusStateChanges = true;
+    node.getFocusableTree().onTreeFocus(node, prevTree);
+    node.onNodeFocus();
+    this.lockFocusStateChanges = false;
+
+    this.setNodeToVisualActiveFocus(node);
+    node.getFocusableElement().focus();
+  }
+
+  /**
+   * Marks the specified node as passively focused, also calling related
+   * lifecycle callback methods for both the node and its parent tree. This
+   * ensures that the node is properly styled to indicate its passive focus.
+   *
+   * This does not change the manager's currently tracked node, nor does it
+   * change any other nodes.
+   *
+   * @param node The node to be passively focused.
+   * @param nextTree The tree of the node receiving active focus, or null if no
+   *     node will be actively focused.
+   */
+  private passivelyFocusNode(
+    node: IFocusableNode,
+    nextTree: IFocusableTree | null,
+  ): void {
+    this.lockFocusStateChanges = true;
+    node.getFocusableTree().onTreeBlur(nextTree);
+    node.onNodeBlur();
+    this.lockFocusStateChanges = false;
+
+    this.setNodeToVisualPassiveFocus(node);
+  }
+
+  /**
+   * Updates the node's styling to indicate that it should have an active focus
+   * indicator.
+   *
+   * @param node The node to be styled for active focus.
+   */
+  private setNodeToVisualActiveFocus(node: IFocusableNode): void {
     const element = node.getFocusableElement();
     dom.addClass(element, FocusManager.ACTIVE_FOCUS_NODE_CSS_CLASS_NAME);
     dom.removeClass(element, FocusManager.PASSIVE_FOCUS_NODE_CSS_CLASS_NAME);
-    element.focus();
   }
 
-  private setNodeToPassive(node: IFocusableNode): void {
+  /**
+   * Updates the node's styling to indicate that it should have a passive focus
+   * indicator.
+   *
+   * @param node The node to be styled for passive focus.
+   */
+  private setNodeToVisualPassiveFocus(node: IFocusableNode): void {
     const element = node.getFocusableElement();
     dom.removeClass(element, FocusManager.ACTIVE_FOCUS_NODE_CSS_CLASS_NAME);
     dom.addClass(element, FocusManager.PASSIVE_FOCUS_NODE_CSS_CLASS_NAME);
   }
 
+  /**
+   * Removes any active/passive indicators for the specified node.
+   *
+   * @param node The node which should have neither passive nor active focus
+   *     indication.
+   */
   private removeHighlight(node: IFocusableNode): void {
     const element = node.getFocusableElement();
     dom.removeClass(element, FocusManager.ACTIVE_FOCUS_NODE_CSS_CLASS_NAME);

@@ -25,8 +25,14 @@ export interface IFocusableNode {
    * and a tab index must be present in order for the element to be focusable in
    * the DOM).
    *
-   * It's expected the return element will not change for the lifetime of the
-   * node.
+   * The returned element must be visible if the node is ever focused via
+   * FocusManager.focusNode() or FocusManager.focusTree(). It's allowed for an
+   * element to be hidden until onNodeFocus() is called, or become hidden with a
+   * call to onNodeBlur().
+   *
+   * It's expected the actual returned element will not change for the lifetime
+   * of the node (that is, its properties can change but a new element should
+   * never be returned).
    */
   getFocusableElement(): HTMLElement | SVGElement;
 
@@ -36,4 +42,38 @@ export interface IFocusableNode {
    * belongs.
    */
   getFocusableTree(): IFocusableTree;
+
+  /**
+   * Called when this node receives active focus.
+   *
+   * Note that it's fine for implementations to change visibility modifiers, but
+   * they should avoid the following:
+   * - Creating or removing DOM elements (including via the renderer or drawer).
+   * - Affecting focus via DOM focus() calls or the FocusManager.
+   */
+  onNodeFocus(): void;
+
+  /**
+   * Called when this node loses active focus. It may still have passive focus.
+   *
+   * This has the same implementation restrictions as onNodeFocus().
+   */
+  onNodeBlur(): void;
+}
+
+/**
+ * Determines whether the provided object fulfills the contract of
+ * IFocusableNode.
+ *
+ * @param object The object to test.
+ * @returns Whether the provided object can be used as an IFocusableNode.
+ */
+export function isFocusableNode(object: any | null): object is IFocusableNode {
+  return (
+    object &&
+    'getFocusableElement' in object &&
+    'getFocusableTree' in object &&
+    'onNodeFocus' in object &&
+    'onNodeBlur' in object
+  );
 }