From 49844e8baa069cf249758779d785c09bbdbae6af Mon Sep 17 00:00:00 2001
From: sakit0 <sakito21@gmail.com>
Date: Mon, 16 Dec 2024 20:09:42 +0900
Subject: [PATCH] feat: pass agentId as a parameter to the server action

---
 app/(playground)/p/[agentId]/canary/contexts/graph.tsx | 9 ++++++---
 app/(playground)/p/[agentId]/canary/page.tsx           | 5 +++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/app/(playground)/p/[agentId]/canary/contexts/graph.tsx b/app/(playground)/p/[agentId]/canary/contexts/graph.tsx
index 89b0f2e9..1c54a391 100644
--- a/app/(playground)/p/[agentId]/canary/contexts/graph.tsx
+++ b/app/(playground)/p/[agentId]/canary/contexts/graph.tsx
@@ -11,6 +11,7 @@ import {
 } from "react";
 import { deriveFlows } from "../lib/graph";
 import type {
+	AgentId,
 	Artifact,
 	Connection,
 	ConnectionId,
@@ -223,6 +224,7 @@ export function GraphContextProvider({
 	defaultGraph,
 	defaultGraphUrl,
 	onPersistAction,
+	agentId,
 }: {
 	children: ReactNode;
 	defaultGraph: Graph;
@@ -231,7 +233,8 @@ export function GraphContextProvider({
 	 * Persists the graph to the server.
 	 * Returns the new graph URL.
 	 */
-	onPersistAction: (graph: Graph) => Promise<string>;
+	onPersistAction: (agentId: AgentId, graph: Graph) => Promise<string>;
+	agentId: AgentId;
 }) {
 	const graphRef = useRef(defaultGraph);
 	const [graph, setGraph] = useState(graphRef.current);
@@ -241,14 +244,14 @@ export function GraphContextProvider({
 	const persist = useCallback(async () => {
 		isPendingPersistRef.current = false;
 		try {
-			const newGraphUrl = await onPersistAction(graphRef.current);
+			const newGraphUrl = await onPersistAction(agentId, graphRef.current);
 			setGraphUrl(newGraphUrl);
 			return newGraphUrl;
 		} catch (error) {
 			console.error("Failed to persist graph:", error);
 			return graphUrl;
 		}
-	}, [onPersistAction, graphUrl]);
+	}, [onPersistAction, graphUrl, agentId]);
 
 	const flush = useCallback(async () => {
 		if (persistTimeoutRef.current) {
diff --git a/app/(playground)/p/[agentId]/canary/page.tsx b/app/(playground)/p/[agentId]/canary/page.tsx
index 5220a1f7..83dc8766 100644
--- a/app/(playground)/p/[agentId]/canary/page.tsx
+++ b/app/(playground)/p/[agentId]/canary/page.tsx
@@ -60,7 +60,7 @@ export default async function Page({
 		(res) => res.json() as unknown as Graph,
 	);
 
-	async function persistGraph(graph: Graph) {
+	async function persistGraph(agentId: AgentId, graph: Graph) {
 		"use server";
 		const { url } = await putGraph(graph);
 		await db
@@ -85,7 +85,7 @@ export default async function Page({
 	let graphUrl = agent.graphUrl;
 	if (!isLatestVersion(graph)) {
 		graph = migrateGraph(graph);
-		graphUrl = await persistGraph(graph);
+		graphUrl = await persistGraph(agentId, graph);
 	}
 
 	async function updateAgentName(agentName: string) {
@@ -131,6 +131,7 @@ export default async function Page({
 				defaultGraph={graph}
 				onPersistAction={persistGraph}
 				defaultGraphUrl={graphUrl}
+				agentId={agentId}
 			>
 				<PropertiesPanelProvider>
 					<ReactFlowProvider>