diff --git a/app/page.tsx b/app/page.tsx
index ad773a9..9dd92fb 100644
--- a/app/page.tsx
+++ b/app/page.tsx
@@ -1,83 +1,44 @@
 "use client";
-import { DrawIoEmbed, DrawIoEmbedRef } from "react-drawio";
-
-import { useRef, useState } from "react";
-import { extractDiagramXML } from "./extract_xml";
+import React from "react";
+import { DrawIoEmbed } from "react-drawio";
 import ChatPanel from "@/components/chat-panel";
+import { DiagramProvider, useDiagram } from "@/contexts/diagram-context";
 
-export default function Home() {
-    const drawioRef = useRef<DrawIoEmbedRef>(null);
-    const [chartXML, setChartXML] = useState<string>("");
-    // Add a ref to store the resolver function
-    const resolverRef = useRef<((value: string) => void) | null>(null);
-    // Add state for diagram history
-    const [diagramHistory, setDiagramHistory] = useState<
-        { svg: string; xml: string }[]
-    >([]);
-    // Add state for latest SVG
-    const [latestSvg, setLatestSvg] = useState<string>("");
-
-    const handleExport = () => {
-        if (drawioRef.current) {
-            drawioRef.current.exportDiagram({
-                format: "xmlsvg",
-            });
-        }
-    };
-
-    const loadDiagram = (chart: string) => {
-        if (drawioRef.current) {
-            drawioRef.current.load({
-                xml: chart,
-            });
-        }
-    };
+// Internal layout component
+function DiagramPageLayout({ children }: { children: React.ReactNode }) {
+    const { drawioRef, handleDiagramExport } = useDiagram();
 
     return (
         <div className="flex h-screen bg-gray-100">
-            <div className="w-2/3 p-1">
-                <DrawIoEmbed
-                    ref={drawioRef}
-                    onExport={(data) => {
-                        const extractedXML = extractDiagramXML(data.data);
-                        setChartXML(extractedXML);
-                        // Store the latest SVG data
-                        setLatestSvg(data.data);
-                        // Directly update diagramHistory with the new data
-                        setDiagramHistory((prev) => [
-                            ...prev,
-                            { svg: data.data, xml: extractedXML },
-                        ]);
-                        // If there's a pending resolver, resolve it with the fresh XML
-                        if (resolverRef.current) {
-                            resolverRef.current(extractedXML);
-                            resolverRef.current = null;
-                        }
-                    }}
-                    urlParameters={{
-                        spin: true,
-                        libraries: false,
-                        saveAndExit: false,
-                        noExitBtn: true,
-                    }}
-                />
-            </div>
-            <div className="w-1/3 p-1 border-gray-300">
-                <ChatPanel
-                    chartXML={chartXML}
-                    onDisplayChart={(xml) => loadDiagram(xml)}
-                    onFetchChart={() => {
-                        return new Promise<string>((resolve) => {
-                            // Store the resolver so onExport can use it
-                            resolverRef.current = resolve;
-                            // Trigger the export
-                            handleExport();
-                        });
-                    }}
-                    diagramHistory={diagramHistory}
-                    onAddToHistory={() => {}}
-                />
+            <div className="w-2/3 p-1 h-full">
+                <div className="h-full relative">
+                    <div className="absolute inset-0">
+                        <div className="w-full h-full">
+                            <DrawIoEmbed
+                                ref={drawioRef}
+                                onExport={handleDiagramExport}
+                                urlParameters={{
+                                    spin: true,
+                                    libraries: false,
+                                    saveAndExit: false,
+                                    noExitBtn: true,
+                                }}
+                            />
+                        </div>
+                    </div>
+                </div>
             </div>
+            <div className="w-1/3 h-full p-1">{children}</div>
         </div>
     );
 }
+
+export default function Home() {
+    return (
+        <DiagramProvider>
+            <DiagramPageLayout>
+                <ChatPanel />
+            </DiagramPageLayout>
+        </DiagramProvider>
+    );
+}
diff --git a/components/chat-input.tsx b/components/chat-input.tsx
index db6503b..01ebbf5 100644
--- a/components/chat-input.tsx
+++ b/components/chat-input.tsx
@@ -27,17 +27,16 @@ import {
 } from "@/components/ui/dialog";
 import Image from "next/image";
 
+import { useDiagram } from "@/contexts/diagram-context";
+
 interface ChatInputProps {
     input: string;
     status: "submitted" | "streaming" | "ready" | "error";
     onSubmit: (e: React.FormEvent<HTMLFormElement>) => void;
     onChange: (e: React.ChangeEvent<HTMLTextAreaElement>) => void;
     setMessages: (messages: any[]) => void;
-    onDisplayChart: (xml: string) => void;
     files?: FileList;
     onFileChange?: (files: FileList | undefined) => void;
-    diagramHistory?: { svg: string; xml: string }[];
-    onSelectHistoryItem?: (xml: string) => void;
     showHistory?: boolean;
     setShowHistory?: (show: boolean) => void;
 }
@@ -48,14 +47,12 @@ export function ChatInput({
     onSubmit,
     onChange,
     setMessages,
-    onDisplayChart,
     files,
     onFileChange,
-    diagramHistory = [],
-    onSelectHistoryItem = () => {},
     showHistory = false,
     setShowHistory = () => {},
 }: ChatInputProps) {
+    const { loadDiagram: onDisplayChart, diagramHistory } = useDiagram();
     const textareaRef = useRef<HTMLTextAreaElement>(null);
     const fileInputRef = useRef<HTMLInputElement>(null);
     const [isDragging, setIsDragging] = useState(false);
@@ -289,9 +286,10 @@ export function ChatInput({
                                         <div
                                             key={index}
                                             className="border rounded-md p-2 cursor-pointer hover:border-primary transition-colors"
-                                            onClick={() =>
-                                                onSelectHistoryItem(item.xml)
-                                            }
+                                            onClick={() => {
+                                                onDisplayChart(item.xml);
+                                                setShowHistory(false);
+                                            }}
                                         >
                                             <div className="aspect-video bg-white rounded overflow-hidden flex items-center justify-center">
                                                 <Image
diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx
index e278cb1..9067b7b 100644
--- a/components/chat-message-display.tsx
+++ b/components/chat-message-display.tsx
@@ -8,23 +8,22 @@ import ExamplePanel from "./chat-example-panel";
 import { Message } from "ai";
 import { convertToLegalXml, replaceNodes } from "@/lib/utils";
 
+import { useDiagram } from "@/contexts/diagram-context";
+
 interface ChatMessageDisplayProps {
-    chartXML: string;
     messages: Message[];
     error?: Error | null;
     setInput: (input: string) => void;
     setFiles: (files: FileList | undefined) => void;
-    onDisplayChart: (xml: string) => void;
 }
 
 export function ChatMessageDisplay({
-    chartXML,
     messages,
     error,
     setInput,
     setFiles,
-    onDisplayChart,
 }: ChatMessageDisplayProps) {
+    const { chartXML, loadDiagram: onDisplayChart } = useDiagram();
     const messagesEndRef = useRef<HTMLDivElement>(null);
     const previousXML = useRef<string>("");
     useEffect(() => {
diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx
index c71e6e2..f707632 100644
--- a/components/chat-panel.tsx
+++ b/components/chat-panel.tsx
@@ -13,21 +13,23 @@ import {
 import { useChat } from "@ai-sdk/react";
 import { ChatInput } from "@/components/chat-input";
 import { ChatMessageDisplay } from "./chat-message-display";
-interface ChatPanelProps {
-    chartXML: string;
-    onDisplayChart: (xml: string) => void;
-    onFetchChart: () => Promise<string>;
-    diagramHistory?: { svg: string; xml: string }[];
-    onAddToHistory?: () => void;
-}
+import { useDiagram } from "@/contexts/diagram-context";
 
-export default function ChatPanel({
-    chartXML,
-    onDisplayChart,
-    onFetchChart,
-    diagramHistory = [],
-    onAddToHistory = () => {},
-}: ChatPanelProps) {
+export default function ChatPanel() {
+    const {
+        chartXML,
+        loadDiagram: onDisplayChart,
+        handleExport: onExport,
+        resolverRef,
+        diagramHistory,
+    } = useDiagram();
+
+    const onFetchChart = () => {
+        return new Promise<string>((resolve) => {
+            resolverRef.current = resolve; // Store the resolver
+            onExport(); // Trigger the export
+        });
+    };
     // Add a step counter to track updates
 
     // Add state for file attachments
@@ -106,12 +108,10 @@ export default function ChatPanel({
             </CardHeader>
             <CardContent className="flex-grow overflow-hidden px-2">
                 <ChatMessageDisplay
-                    chartXML={chartXML}
                     messages={messages}
                     error={error}
                     setInput={setInput}
                     setFiles={handleFileChange}
-                    onDisplayChart={onDisplayChart}
                 />
             </CardContent>
 
@@ -122,11 +122,8 @@ export default function ChatPanel({
                     onSubmit={onFormSubmit}
                     onChange={handleInputChange}
                     setMessages={setMessages}
-                    onDisplayChart={onDisplayChart}
                     files={files}
                     onFileChange={handleFileChange}
-                    diagramHistory={diagramHistory}
-                    onSelectHistoryItem={handleSelectHistoryItem}
                     showHistory={showHistory}
                     setShowHistory={setShowHistory}
                 />
diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx
new file mode 100644
index 0000000..adce5fa
--- /dev/null
+++ b/contexts/diagram-context.tsx
@@ -0,0 +1,86 @@
+"use client";
+
+import React, { createContext, useContext, useRef, useState } from "react";
+import type { DrawIoEmbedRef } from "react-drawio";
+import { extractDiagramXML } from "@/app/extract_xml";
+
+interface DiagramContextType {
+    chartXML: string;
+    latestSvg: string;
+    diagramHistory: { svg: string; xml: string }[];
+    loadDiagram: (chart: string) => void;
+    handleExport: () => void;
+    resolverRef: React.MutableRefObject<((value: string) => void) | null>;
+    drawioRef: React.MutableRefObject<DrawIoEmbedRef | null>;
+    handleDiagramExport: (data: any) => void;
+}
+
+const DiagramContext = createContext<DiagramContextType | undefined>(undefined);
+
+export function DiagramProvider({ children }: { children: React.ReactNode }) {
+    const [chartXML, setChartXML] = useState<string>("");
+    const [latestSvg, setLatestSvg] = useState<string>("");
+    const [diagramHistory, setDiagramHistory] = useState<
+        { svg: string; xml: string }[]
+    >([]);
+    const drawioRef = useRef<DrawIoEmbedRef>(null);
+    const resolverRef = useRef<((value: string) => void) | null>(null);
+
+    const handleExport = () => {
+        if (drawioRef.current) {
+            drawioRef.current.exportDiagram({
+                format: "xmlsvg",
+            });
+        }
+    };
+
+    const loadDiagram = (chart: string) => {
+        if (drawioRef.current) {
+            drawioRef.current.load({
+                xml: chart,
+            });
+        }
+    };
+
+    const handleDiagramExport = (data: any) => {
+        const extractedXML = extractDiagramXML(data.data);
+        setChartXML(extractedXML);
+        setLatestSvg(data.data);
+        setDiagramHistory((prev) => [
+            ...prev,
+            {
+                svg: data.data,
+                xml: extractedXML,
+            },
+        ]);
+        if (resolverRef.current) {
+            resolverRef.current(extractedXML);
+            resolverRef.current = null;
+        }
+    };
+
+    return (
+        <DiagramContext.Provider
+            value={{
+                chartXML,
+                latestSvg,
+                diagramHistory,
+                loadDiagram,
+                handleExport,
+                resolverRef,
+                drawioRef,
+                handleDiagramExport,
+            }}
+        >
+            {children}
+        </DiagramContext.Provider>
+    );
+}
+
+export function useDiagram() {
+    const context = useContext(DiagramContext);
+    if (context === undefined) {
+        throw new Error("useDiagram must be used within a DiagramProvider");
+    }
+    return context;
+}