Support MergedRef

packages/hyperion-autologging/src/ALSurface.ts

@@ -13,7 +13,7 @@ import * as IReactElementVisitor from 'hyperion-react/src/IReactElementVisitor';
 import * as IReactFlowlet from "hyperion-react/src/IReactFlowlet";
 import * as IReactPropsExtension from "hyperion-react/src/IReactPropsExtension";
 import * as Types from "hyperion-util/src/Types";
-import type * as React from 'react';
+import * as React from 'react';
 import { ALFlowletDataType, IALFlowlet } from "./ALFlowletManager";
 import { AUTO_LOGGING_NON_INTERACTIVE_SURFACE, AUTO_LOGGING_SURFACE, SURFACE_SEPARATOR, SURFACE_WRAPPER_ATTRIBUTE_NAME } from './ALSurfaceConsts';
 import * as ALSurfaceContext from "./ALSurfaceContext";
@@ -66,7 +66,7 @@ export type ALSurfaceProps = Readonly<{
   metadata?: ALMetadataEvent['metadata'];
   uiEventMetadata?: EventMetadata,
   capability?: ALSurfaceCapability,
-  nodeRef?: React.RefObject<HTMLElement | null | undefined>,
+  nodeRef?: React.Ref<HTMLElement | null | undefined>,
 }>;
 
 export type ALSurfaceRenderer = (node: React.ReactNode) => React.ReactElement;
@@ -236,6 +236,17 @@ function setupDomElementSurfaceAttribute(options: InitOptions): void {
   });
 }
 
+function getRefElement<T>(ref: React.Ref<T> | undefined | null): T | null {
+  if (!ref) return null;
+  if (typeof ref === 'function') {
+    // Function refs don't expose current; you can't get element here safely
+    // unless you track it yourself, so return null
+    return null;
+  } else if ('current' in ref) {
+    return ref.current;
+  }
+  return null;
+}
 
 export function init(options: InitOptions): ALSurfaceHOC {
   const { flowletManager, channel } = options;
@@ -259,7 +270,7 @@ export function init(options: InitOptions): ALSurfaceHOC {
     const { surface: parentSurface, nonInteractiveSurface: parentNonInteractiveSurface } = surfaceCtx;
 
     let addSurfaceWrapper = props.nodeRef == null;
-    let localRef = ReactModule.useRef<Element>();
+    const localRef = ReactModule.useRef<Element | null>(null);
 
     // empty .capability field is default, means all enabled!
     const capability = props.capability ?? proxiedContext?.mainContext.capability;
@@ -323,6 +334,16 @@ export function init(options: InitOptions): ALSurfaceHOC {
 
     const isProxy = proxiedContext != null;
 
+    const mergedRef = (node: Element | null) => {
+      const htmlNode = node instanceof HTMLElement ? node : null;
+      if (typeof props.nodeRef === 'function') {
+        props.nodeRef(htmlNode);
+      } else if (props.nodeRef && 'current' in props.nodeRef) {
+        (props.nodeRef as React.MutableRefObject<Element | null>).current = node;
+      }
+      localRef.current = node;
+    };
+
     metadata.original_call_flowlet = callFlowlet.getFullName();
     metadata.surface_capability = surfaceCapabilityToString(capability);
     // Update the metadata on every render to ensure it stays current
@@ -342,9 +363,10 @@ export function init(options: InitOptions): ALSurfaceHOC {
      */
     const nodeRef = props.nodeRef ?? localRef;
 
+
     ReactModule.useLayoutEffect(() => {
       __DEV__ && assert(nodeRef != null, "Invalid surface effect without a ref: " + surface);
-      const element = nodeRef.current;
+      const element = getRefElement(nodeRef as React.Ref<HTMLElement>);
       if (element == null) {
         return;
       }
@@ -435,7 +457,7 @@ export function init(options: InitOptions): ALSurfaceHOC {
             [SURFACE_WRAPPER_ATTRIBUTE_NAME]: "1",
             style: { display: 'contents' },
             [domAttributeName]: domAttributeValue,
-            ref: localRef, // addSurfaceWrapper would have been false if a rep was passed in props
+            ref: mergedRef, // addSurfaceWrapper would have been false if a rep was passed in props
           },
           props.children
         );

packages/hyperion-react-testapp/src/App.tsx

@@ -18,6 +18,7 @@ import { PortalBodyContainerComponent } from './component/PortalComponent';
 import TextComponent from './component/TextComponent';
 import RecursiveFuncComponent from './component/RecursiveFuncComponent';
 import SVGClickComponent from './component/SVGClickComponent';
+import { RefNodeComponent } from './component/RefNodeComponent';
 
 function InitComp() {
   const [count, setCount] = React.useState(0);
@@ -60,6 +61,9 @@ const Modes = {
     <ElementNameComponent />
     <TextComponent />
   </div>,
+  'RefNode': () => <div>
+    <RefNodeComponent />
+  </div>,
 };
 type ModeNames = keyof typeof Modes;
 const PersistedOptionValue = new LocalStoragePersistentData<ModeNames>(

packages/hyperion-react-testapp/src/component/RefNodeComponent.tsx

@@ -0,0 +1,71 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved.
+ */
+
+import React, { useEffect, useRef, useCallback, useState } from 'react';
+import { SurfaceComp } from './Surface';
+
+// Utility to merge multiple refs
+function mergeRefs<T>(...refs: (React.Ref<T> | undefined)[]): React.RefCallback<T> {
+  return (node: T) => {
+    refs.forEach(ref => {
+      if (typeof ref === 'function') {
+        ref(node);
+      } else if (ref && 'current' in ref) {
+        (ref as React.MutableRefObject<T | null>).current = node;
+      }
+    });
+  };
+}
+
+const ComponenetWithRefs: React.FC<{ externalRef?: React.Ref<HTMLElement> }> = ({ externalRef }) => {
+  const internalRef = useRef<HTMLElement | null>(null);
+
+  const combinedRef = mergeRefs<HTMLElement>(externalRef, internalRef);
+
+  useEffect(() => {
+    if (internalRef.current) {
+      internalRef.current.style.border = '2px solid blue';
+      internalRef.current.textContent = 'Merged using useMergeRefs!';
+    }
+  }, []);
+
+  return (
+    <div
+      ref={combinedRef}
+      style={{ padding: '12px', marginTop: '20px', backgroundColor: '#eee' }}
+    >
+    A box with merged refs.
+    </div>
+  );
+};
+
+// Component that uses a merged ref
+type RefNodeComponentProps = {
+  externalRef?: React.Ref<HTMLElement>;
+};
+
+export const RefNodeComponent: React.FC<RefNodeComponentProps> = () => {
+  const [element, setElement] = useState<HTMLElement | null>(null);
+
+  // External callback ref
+  const externalRef = (node: HTMLElement | null) => {
+    console.log('External ref called with:', node);
+    setElement(node);
+  };
+
+  return (
+    <SurfaceComp surface='nodeRefs'>
+      <div style={{ padding: '20px' }}>
+        <h3>React Merged Ref Example</h3>
+        <ComponenetWithRefs externalRef={externalRef} />
+        {element && (
+          <p>
+            The tag name of the externally tracked element is: <strong>{element.tagName}</strong>
+          </p>
+        )}
+      </div>
+    </SurfaceComp>
+  );
+
+};