🐛 fix: Enhance Agent Duplication Logic and Filter Sensitive Data

danny-avila · Dec 17, 2024 · 880c4a0 · 880c4a0
1 parent 5252fd3
commit 880c4a0
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 47 deletions.
diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js
@@ -189,18 +189,24 @@ const updateAgentHandler = async (req, res) => {
  * @returns {Agent} 201 - success response - application/json
  */
 const duplicateAgentHandler = async (req, res) => {
-  try {
-    const id = req.params.id;
-    const { id: userId } = req.user;
+  const { id } = req.params;
+  const { id: userId } = req.user;
 
+  try {
     const agent = await getAgent({ id });
     if (!agent) {
-      return res.status(404).json({ error: 'Agent not found' });
+      return res.status(404).json({
+        error: 'Agent not found',
+        status: 'error',
+      });
     }
 
     const { name, description, instructions, tools, provider, model } = agent;
-    const newAgent = await createAgent({
-      id: `agent_${nanoid()}`,
+
+    const newAgentId = `agent_${nanoid()}`;
+
+    await createAgent({
+      id: newAgentId,
       author: userId,
       name: `${name} (Copy)`,
       description,
@@ -211,46 +217,65 @@ const duplicateAgentHandler = async (req, res) => {
       actions: [],
     });
 
+    const actionsData = [];
     const originalActions = await getActions({ agent_id: id }, true);
-    let actionsData = [];
-
-    if (originalActions && originalActions.length > 0) {
-      const newActions = [];
-
-      for (const action of originalActions) {
-        const newActionId = nanoid();
-        const [domain] = action.action_id.split(actionDelimiter);
-
-        const fullActionId = `${domain}${actionDelimiter}${newActionId}`;
-
-        await updateAction(
-          { action_id: newActionId },
-          {
-            metadata: action.metadata,
-            agent_id: newAgent.id,
-            user: userId,
-          },
-        );
-
-        newActions.push(fullActionId);
-        actionsData.push({
-          id: newActionId,
-          action_id: fullActionId,
-          metadata: action.metadata,
-          domain,
-        });
-      }
 
-      await updateAgent({ id: newAgent.id }, { actions: newActions });
+    const sensitiveFields = ['api_key', 'oauth_client_id', 'oauth_client_secret'];
+
+    if (originalActions?.length) {
+      const newActions = await Promise.all(
+        originalActions.map(async (action) => {
+          const newActionId = nanoid();
+          const [domain] = action.action_id.split(actionDelimiter);
+          const fullActionId = `${domain}${actionDelimiter}${newActionId}`;
+
+          const filteredMetadata = { ...action.metadata };
+          for (const field of sensitiveFields) {
+            delete filteredMetadata[field];
+          }
+
+          await updateAction(
+            { action_id: newActionId },
+            {
+              metadata: filteredMetadata,
+              agent_id: newAgentId,
+              user: userId,
+            },
+          );
+
+          actionsData.push({
+            id: newActionId,
+            action_id: fullActionId,
+            metadata: filteredMetadata,
+            domain,
+          });
+
+          return fullActionId;
+        }),
+      );
+
+      await updateAgent({ id: newAgentId }, { actions: newActions });
     }
 
-    const finalAgent = await getAgent({ id: newAgent.id });
+    const finalAgent = await getAgent({ id: newAgentId });
+
+    const filteredActions = actionsData.map((action) => {
+      const filteredMetadata = { ...action.metadata };
+      for (const field of sensitiveFields) {
+        delete filteredMetadata[field];
+      }
+      return { ...action, metadata: filteredMetadata };
+    });
+
     return res.status(201).json({
       agent: finalAgent,
-      actions: actionsData,
+      actions: filteredActions,
     });
   } catch (error) {
-    logger.error('[/Agents/:id/duplicate] Error duplicating Agent', error);
+    logger.error('[/Agents/:id/duplicate] Error duplicating Agent:', {
+      error,
+    });
+
     res.status(500).json({ error: error.message });
   }
 };

diff --git a/client/src/components/SidePanel/Agents/AgentSelect.tsx b/client/src/components/SidePanel/Agents/AgentSelect.tsx
@@ -186,7 +186,7 @@ export default function AgentSelect({
       )}
       className={cn(
         'rounded-md dark:border-gray-700 dark:bg-gray-850',
-        'max-w z-50 flex h-[40px] w-full flex-none items-center justify-center truncate px-4 hover:cursor-pointer hover:border-green-500 focus:border-gray-400',
+        'z-50 flex h-[40px] w-full flex-none items-center justify-center truncate px-4 hover:cursor-pointer hover:border-green-500 focus:border-gray-400',
       )}
       renderOption={() => (
         <span className="flex items-center gap-1.5 truncate">

diff --git a/client/src/data-provider/Agents/mutations.ts b/client/src/data-provider/Agents/mutations.ts
@@ -149,15 +149,16 @@ export const useDuplicateAgentMutation = (
         }
 
         const existingActions = queryClient.getQueryData<t.Action[]>([QueryKeys.actions]) || [];
-        const actionsMap = new Map(existingActions.map((a) => [a.action_id, a]));
 
-        actions.forEach((action) => {
-          const { assistant_id, ...actionWithoutAssistant } = action;
-          const updatedAction = { ...actionWithoutAssistant, agent_id: agent.id };
-          actionsMap.set(updatedAction.action_id, updatedAction);
-        });
-
-        queryClient.setQueryData<t.Action[]>([QueryKeys.actions], Array.from(actionsMap.values()));
+        queryClient.setQueryData<t.Action[]>(
+          [QueryKeys.actions],
+          existingActions.concat(
+            actions.map(({ assistant_id, ...action }) => ({
+              ...action,
+              agent_id: agent.id,
+            })),
+          ),
+        );
 
         return options?.onSuccess?.({ agent, actions }, variables, context);
       },