diff --git a/src/core/client/ui/components/Button/Button.css b/src/core/client/ui/components/Button/Button.css
index e3cab49f2..3a0151ecc 100644
--- a/src/core/client/ui/components/Button/Button.css
+++ b/src/core/client/ui/components/Button/Button.css
@@ -2,7 +2,7 @@
   border-radius: var(--round-corners);
   background-color: transparent;
   position: relative;
-  display: flex;
+  display: inline-flex;
   justify-content: center;
   align-items: center;
   box-sizing: border-box;
diff --git a/src/core/client/ui/components/TrapFocus/TrapFocus.mdx b/src/core/client/ui/components/TrapFocus/TrapFocus.mdx
index d5c209aa9..da30e1171 100644
--- a/src/core/client/ui/components/TrapFocus/TrapFocus.mdx
+++ b/src/core/client/ui/components/TrapFocus/TrapFocus.mdx
@@ -5,6 +5,7 @@ menu: UI Kit
 
 import { Playground } from 'docz'
 import TrapFocus from './TrapFocus'
+import Button from '../Button'
 
 # TrapFocus
 
@@ -12,28 +13,15 @@ Traps focus inside component when using keyboard navigation for accessibility pu
 
 ## Basic usage
 
-```ts
-import React from "react";
-
-class Dialog extends React.Component {
-  private firstFocusable: HTMLElement;
-  private lastFocusable: HTMLElement;
-
-  private setFirstFocusable = (ref: HTMLElement) => (this.firstFocusable = ref);
-  private setLastFocusable = (ref: HTMLElement) => (this.lastFocusable = ref);
-
-  public render() {
-    return (
+<Playground>
+  <div>
+    <TrapFocus>
+    {({firstFocusableRef, lastFocusableRef}) =>
       <div>
-        <TrapFocus
-          firstFocusable={this.firstFocusable}
-          lastFocusable={this.lastFocusable}
-        >
-          <TextField ref={this.setFirstFocusable} />
-          <Button ref={this.setLastFocusable}>Send</Button>
-        </TrapFocus>
+        <Button forwardRef={firstFocusableRef} autoFocus>Cancel</Button>
+        <Button forwardRef={lastFocusableRef}>Send</Button>
       </div>
-    );
-  }
-}
-```
+    }
+    </TrapFocus>
+  </div>
+</Playground>
diff --git a/src/core/client/ui/components/TrapFocus/TrapFocus.spec.tsx b/src/core/client/ui/components/TrapFocus/TrapFocus.spec.tsx
index f4932775f..3abcdb693 100644
--- a/src/core/client/ui/components/TrapFocus/TrapFocus.spec.tsx
+++ b/src/core/client/ui/components/TrapFocus/TrapFocus.spec.tsx
@@ -5,11 +5,16 @@ import Sinon from "sinon";
 import { PropTypesOf } from "talk-ui/types";
 import TrapFocus from "./TrapFocus";
 
+const FakeFocusable: any = class extends React.Component {
+  public focus = Sinon.spy();
+  public render() {
+    return null;
+  }
+};
+
 it("renders correctly", () => {
   const props: PropTypesOf<TrapFocus> = {
-    firstFocusable: null,
-    lastFocusable: null,
-    children: (
+    children: () => (
       <>
         <span>child1</span>
         <span>child2</span>
@@ -25,15 +30,13 @@ it("renders correctly", () => {
 });
 
 it("Change focus to `lastFocusable` when focus reaches beginning", () => {
-  const fakeHTMLElementBegin = { focus: Sinon.spy() };
-  const fakeHTMLElementEnd = { focus: Sinon.spy() };
   const props: PropTypesOf<TrapFocus> = {
-    firstFocusable: fakeHTMLElementBegin as any,
-    lastFocusable: fakeHTMLElementEnd as any,
-    children: (
+    children: ({ firstFocusableRef, lastFocusableRef }) => (
       <>
+        <FakeFocusable ref={firstFocusableRef} />
         <span>child1</span>
         <span>child2</span>
+        <FakeFocusable ref={lastFocusableRef} />
       </>
     ),
   };
@@ -43,20 +46,22 @@ it("Change focus to `lastFocusable` when focus reaches beginning", () => {
     </div>
   );
   renderer.root.findAllByProps({ tabIndex: 0 })[0].props.onFocus();
-  expect(fakeHTMLElementBegin.focus.called).toBe(false);
-  expect(fakeHTMLElementEnd.focus.called).toBe(true);
+  expect(
+    renderer.root.findAllByType(FakeFocusable)[0].instance.focus.called
+  ).toBe(false);
+  expect(
+    renderer.root.findAllByType(FakeFocusable)[1].instance.focus.called
+  ).toBe(true);
 });
 
-it("Change focus to `firstFocusable` when focus reaches the end", () => {
-  const fakeHTMLElementBegin = { focus: Sinon.spy() };
-  const fakeHTMLElementEnd = { focus: Sinon.spy() };
+it("Change focus to `firstFocusable` when focus reaches end", () => {
   const props: PropTypesOf<TrapFocus> = {
-    firstFocusable: fakeHTMLElementBegin as any,
-    lastFocusable: fakeHTMLElementEnd as any,
-    children: (
+    children: ({ firstFocusableRef, lastFocusableRef }) => (
       <>
+        <FakeFocusable ref={firstFocusableRef} />
         <span>child1</span>
         <span>child2</span>
+        <FakeFocusable ref={lastFocusableRef} />
       </>
     ),
   };
@@ -66,6 +71,10 @@ it("Change focus to `firstFocusable` when focus reaches the end", () => {
     </div>
   );
   renderer.root.findAllByProps({ tabIndex: 0 })[1].props.onFocus();
-  expect(fakeHTMLElementBegin.focus.called).toBe(true);
-  expect(fakeHTMLElementEnd.focus.called).toBe(false);
+  expect(
+    renderer.root.findAllByType(FakeFocusable)[0].instance.focus.called
+  ).toBe(true);
+  expect(
+    renderer.root.findAllByType(FakeFocusable)[1].instance.focus.called
+  ).toBe(false);
 });
diff --git a/src/core/client/ui/components/TrapFocus/TrapFocus.tsx b/src/core/client/ui/components/TrapFocus/TrapFocus.tsx
index a50b50d6c..0d6effbbe 100644
--- a/src/core/client/ui/components/TrapFocus/TrapFocus.tsx
+++ b/src/core/client/ui/components/TrapFocus/TrapFocus.tsx
@@ -1,25 +1,34 @@
-import React from "react";
+import React, { RefObject } from "react";
 
 export interface Focusable {
   focus: () => void;
 }
 
+interface RenderProps {
+  firstFocusableRef: RefObject<Focusable>;
+  lastFocusableRef: RefObject<Focusable>;
+}
+
 export interface TrapFocusProps {
-  firstFocusable: Focusable | null;
-  lastFocusable: Focusable | null;
-  children: React.ReactNode;
+  children: (props: RenderProps) => React.ReactNode;
 }
 
 export default class TrapFocus extends React.Component<TrapFocusProps> {
+  private firstFocusableRef = React.createRef<Focusable>();
+  private lastFocusableRef = React.createRef<Focusable>();
+
   // Trap keyboard focus inside the dropdown until a value has been chosen.
-  public focusBegin = () => this.props.firstFocusable!.focus();
-  public focusEnd = () => this.props.lastFocusable!.focus();
+  private focusBegin = () => this.firstFocusableRef.current!.focus();
+  private focusEnd = () => this.lastFocusableRef.current!.focus();
 
   public render() {
     return (
       <>
         <div tabIndex={0} onFocus={this.focusEnd} />
-        {this.props.children}
+        {this.props.children({
+          firstFocusableRef: this.firstFocusableRef,
+          lastFocusableRef: this.lastFocusableRef,
+        })}
         <div tabIndex={0} onFocus={this.focusBegin} />
       </>
     );