fix: pass ref explicitly using iframeProps (#33)

liady · web-flow · commit d01b5d1e4cda · 2025-07-17T00:38:52.000+03:00
diff --git a/packages/client/src/components/HTMLResourceRenderer.tsx b/packages/client/src/components/HTMLResourceRenderer.tsx
@@ -7,15 +7,19 @@ export type HTMLResourceRendererProps = {
   resource: Partial<Resource>;
   onUIAction?: (result: UIActionResult) => Promise<unknown>;
   style?: React.CSSProperties;
-  iframeProps?: Omit<React.HTMLAttributes<HTMLIFrameElement>, 'src' | 'srcDoc' | 'ref' | 'style'>;
+  iframeProps?: Omit<React.HTMLAttributes<HTMLIFrameElement>, 'src' | 'srcDoc' | 'style'> & {
+    ref?: React.RefObject<HTMLIFrameElement>;
+  };
 };
 
-export const HTMLResourceRenderer = React.forwardRef<
-  HTMLIFrameElement | null,
-  HTMLResourceRendererProps
->(({ resource, onUIAction, style, iframeProps }, ref) => {
+export const HTMLResourceRenderer = ({
+  resource,
+  onUIAction,
+  style,
+  iframeProps,
+}: HTMLResourceRendererProps) => {
   const iframeRef = useRef<HTMLIFrameElement | null>(null);
-  useImperativeHandle(ref, () => iframeRef.current as HTMLIFrameElement);
+  useImperativeHandle(iframeProps?.ref, () => iframeRef.current as HTMLIFrameElement);
 
   const { error, iframeSrc, iframeRenderMode, htmlString } = useMemo(
     () => processHTMLResource(resource),
@@ -78,6 +82,6 @@ export const HTMLResourceRenderer = React.forwardRef<
   }
 
   return <p className="text-gray-500">Initializing HTML resource display...</p>;
-});
+};
 
 HTMLResourceRenderer.displayName = 'HTMLResourceRenderer';
diff --git a/packages/client/src/components/UIResourceRenderer.tsx b/packages/client/src/components/UIResourceRenderer.tsx
@@ -1,4 +1,3 @@
-import React from 'react';
 import type { Resource } from '@modelcontextprotocol/sdk/types.js';
 import { ResourceContentType, UIActionResult } from '../types';
 import { HTMLResourceRenderer, HTMLResourceRendererProps } from './HTMLResourceRenderer';
@@ -29,10 +28,7 @@ function getContentType(resource: Partial<Resource>): ResourceContentType | unde
   }
 }
 
-export const UIResourceRenderer = React.forwardRef<
-  HTMLIFrameElement | null,
-  UIResourceRendererProps
->((props, ref) => {
+export const UIResourceRenderer = (props: UIResourceRendererProps) => {
   const { resource, onUIAction, supportedContentTypes, htmlProps, remoteDomProps } = props;
   const contentType = getContentType(resource);
 
@@ -42,15 +38,9 @@ export const UIResourceRenderer = React.forwardRef<
 
   switch (contentType) {
     case 'rawHtml':
-    case 'externalUrl':
-      return (
-        <HTMLResourceRenderer
-          resource={resource}
-          onUIAction={onUIAction}
-          {...htmlProps}
-          ref={ref}
-        />
-      );
+    case 'externalUrl': {
+      return <HTMLResourceRenderer resource={resource} onUIAction={onUIAction} {...htmlProps} />;
+    }
     case 'remoteDom':
       return (
         <RemoteDOMResourceRenderer
@@ -63,6 +53,6 @@ export const UIResourceRenderer = React.forwardRef<
     default:
       return <p className="text-red-500">Unsupported resource type.</p>;
   }
-});
+};
 
 UIResourceRenderer.displayName = 'UIResourceRenderer';
diff --git a/packages/client/src/components/__tests__/HTMLResourceRenderer.test.tsx b/packages/client/src/components/__tests__/HTMLResourceRenderer.test.tsx
@@ -291,7 +291,7 @@ describe('HTMLResource iframe communication', () => {
 
   it('should pass ref to iframe', () => {
     const ref = React.createRef<HTMLIFrameElement>();
-    render(<HTMLResourceRenderer resource={mockResourceBaseForUIActionTests} ref={ref} />);
+    render(<HTMLResourceRenderer resource={mockResourceBaseForUIActionTests} iframeProps={{ ref }} />);
     expect(ref.current).toBeInTheDocument();
   });
 
diff --git a/packages/client/src/components/__tests__/UIResourceRenderer.unmocked.test.tsx b/packages/client/src/components/__tests__/UIResourceRenderer.unmocked.test.tsx
@@ -12,7 +12,7 @@ describe('UIResourceRenderer', () => {
   };
   it('should pass ref to HTMLResourceRenderer', () => {
     const ref = React.createRef<HTMLIFrameElement>();
-    render(<UIResourceRenderer resource={testResource} ref={ref} />);
+    render(<UIResourceRenderer resource={testResource} htmlProps={{ iframeProps: { ref } }} />);
     expect(ref.current).toBeInTheDocument();
   });
 });
