diff --git a/app/(playground)/p/[agentId]/canary/actions.ts b/app/(playground)/p/[agentId]/canary/actions.ts
index 47b942e6..a7881bd8 100644
--- a/app/(playground)/p/[agentId]/canary/actions.ts
+++ b/app/(playground)/p/[agentId]/canary/actions.ts
@@ -1,5 +1,6 @@
 "use server";
 
+import { db } from "@/drizzle";
 import { anthropic } from "@ai-sdk/anthropic";
 import { google } from "@ai-sdk/google";
 import { openai } from "@ai-sdk/openai";
@@ -16,6 +17,7 @@ import * as v from "valibot";
 import { vercelBlobFileFolder, vercelBlobGraphFolder } from "./constants";
 import { textGenerationPrompt } from "./prompts";
 import type {
+	AgentId,
 	ArtifactId,
 	FileData,
 	FileId,
@@ -106,7 +108,7 @@ type ActionSource = TextSource | TextGenerationSource | FileSource;
 
 export async function action(
 	artifactId: ArtifactId,
-	graphUrl: string,
+	agentId: AgentId,
 	nodeId: NodeId,
 ) {
 	const lf = new Langfuse();
@@ -114,7 +116,14 @@ export async function action(
 		sessionId: artifactId,
 	});
 
-	const graph = await fetch(graphUrl).then(
+	const agent = await db.query.agents.findFirst({
+		where: (agents, { eq }) => eq(agents.id, agentId),
+	});
+	if (agent === undefined || agent.graphUrl === null) {
+		throw new Error(`Agent with id ${agentId} not found`);
+	}
+
+	const graph = await fetch(agent.graphUrl).then(
 		(res) => res.json() as unknown as Graph,
 	);
 	const node = graph.nodes.find((node) => node.id === nodeId);
diff --git a/app/(playground)/p/[agentId]/canary/contexts/execution.tsx b/app/(playground)/p/[agentId]/canary/contexts/execution.tsx
index 36195b12..438930f0 100644
--- a/app/(playground)/p/[agentId]/canary/contexts/execution.tsx
+++ b/app/(playground)/p/[agentId]/canary/contexts/execution.tsx
@@ -20,7 +20,6 @@ interface ExecutionProviderProps {
 	children: ReactNode;
 	executeAction: (
 		artifactId: ArtifactId,
-		graphUrl: string,
 		nodeId: NodeId,
 	) => Promise<StreamableValue<TextArtifactObject, unknown>>;
 }
@@ -56,9 +55,9 @@ export function ExecutionProvider({
 				},
 			});
 			setTab("Result");
-			const latestGraphUrl = await flush();
+			await flush();
 			try {
-				const stream = await executeAction(artifactId, latestGraphUrl, nodeId);
+				const stream = await executeAction(artifactId, nodeId);
 
 				let textArtifactObject: TextArtifactObject = {
 					type: "text",
diff --git a/app/(playground)/p/[agentId]/canary/page.tsx b/app/(playground)/p/[agentId]/canary/page.tsx
index d6b3d562..96958ce5 100644
--- a/app/(playground)/p/[agentId]/canary/page.tsx
+++ b/app/(playground)/p/[agentId]/canary/page.tsx
@@ -4,7 +4,6 @@ import { del, list } from "@vercel/blob";
 import { ReactFlowProvider } from "@xyflow/react";
 import { eq } from "drizzle-orm";
 import { notFound } from "next/navigation";
-import type { AgentId } from "../beta-proto/types";
 import { action, putGraph } from "./actions";
 import { Editor } from "./components/editor";
 import { AgentNameProvider } from "./contexts/agent-name";
@@ -16,7 +15,7 @@ import { PropertiesPanelProvider } from "./contexts/properties-panel";
 import { ToastProvider } from "./contexts/toast";
 import { ToolbarContextProvider } from "./contexts/toolbar";
 import { isLatestVersion, migrateGraph } from "./graph";
-import type { ArtifactId, Graph, NodeId } from "./types";
+import type { AgentId, ArtifactId, Graph, NodeId } from "./types";
 import { buildGraphFolderPath } from "./utils";
 
 // Extend the max duration of the server actions from this page to 5 minutes
@@ -88,13 +87,9 @@ export default async function Page({
 		return agentName;
 	}
 
-	async function execute(
-		artifactId: ArtifactId,
-		graphUrl: string,
-		nodeId: NodeId,
-	) {
+	async function execute(artifactId: ArtifactId, nodeId: NodeId) {
 		"use server";
-		return await action(artifactId, graphUrl, nodeId);
+		return await action(artifactId, agentId, nodeId);
 	}
 
 	return (
diff --git a/app/(playground)/p/[agentId]/canary/types.ts b/app/(playground)/p/[agentId]/canary/types.ts
index 2b26c241..3a23a5c4 100644
--- a/app/(playground)/p/[agentId]/canary/types.ts
+++ b/app/(playground)/p/[agentId]/canary/types.ts
@@ -212,3 +212,5 @@ export interface SubGraph {
 	nodes: Set<NodeId>;
 	connections: Set<ConnectionId>;
 }
+
+export type AgentId = `agnt_${string}`;