fix(ui): fix kb id extraction logic for resource, sync tags (#3763)

* fix(ui): fix kb id extraction logic for resource, sync tags

* Pass knowledge base id back on edit tag

---------

Co-authored-by: Theodore Li <theo@sim.ai>

Author: TheodoreSpeaks

apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-registry/resource-registry.tsx

@@ -138,6 +138,7 @@ const RESOURCE_INVALIDATORS: Record<
   knowledgebase: (qc, _wId, id) => {
     qc.invalidateQueries({ queryKey: knowledgeKeys.lists() })
     qc.invalidateQueries({ queryKey: knowledgeKeys.detail(id) })
+    qc.invalidateQueries({ queryKey: knowledgeKeys.tagDefinitions(id) })
   },
 }
 

apps/sim/lib/copilot/resource-extraction.test.ts

@@ -0,0 +1,85 @@
+/**
+ * @vitest-environment node
+ */
+import { describe, expect, it } from 'vitest'
+import { extractResourcesFromToolResult } from './resource-extraction'
+
+describe('extractResourcesFromToolResult', () => {
+  it('uses the knowledge base id for knowledge_base tag mutations', () => {
+    const resources = extractResourcesFromToolResult(
+      'knowledge_base',
+      {
+        operation: 'update_tag',
+        args: {
+          knowledgeBaseId: 'kb_123',
+          tagDefinitionId: 'tag_456',
+        },
+      },
+      {
+        success: true,
+        message: 'Tag updated successfully',
+        data: {
+          id: 'tag_456',
+          displayName: 'Priority',
+          fieldType: 'text',
+        },
+      }
+    )
+
+    expect(resources).toEqual([
+      {
+        type: 'knowledgebase',
+        id: 'kb_123',
+        title: 'Knowledge Base',
+      },
+    ])
+  })
+
+  it('uses knowledgeBaseId from the tool result when update_tag args omit it', () => {
+    const resources = extractResourcesFromToolResult(
+      'knowledge_base',
+      {
+        operation: 'update_tag',
+        args: {
+          tagDefinitionId: 'tag_456',
+        },
+      },
+      {
+        success: true,
+        message: 'Tag updated successfully',
+        data: {
+          id: 'tag_456',
+          knowledgeBaseId: 'kb_123',
+          displayName: 'Priority',
+          fieldType: 'text',
+        },
+      }
+    )
+
+    expect(resources).toEqual([
+      {
+        type: 'knowledgebase',
+        id: 'kb_123',
+        title: 'Knowledge Base',
+      },
+    ])
+  })
+
+  it('does not create resources for read-only knowledge base tag operations', () => {
+    const resources = extractResourcesFromToolResult(
+      'knowledge_base',
+      {
+        operation: 'list_tags',
+        args: {
+          knowledgeBaseId: 'kb_123',
+        },
+      },
+      {
+        success: true,
+        data: [],
+      }
+    )
+
+    expect(resources).toEqual([])
+  })
+})

apps/sim/lib/copilot/resource-extraction.ts

@@ -155,11 +155,13 @@ export function extractResourcesFromToolResult(
     case 'knowledge_base': {
       if (READ_ONLY_KB_OPS.has(getOperation(params) ?? '')) return []
 
+      const args = asRecord(params?.args)
       const kbId =
-        (data.id as string) ??
+        (args.knowledgeBaseId as string) ??
+        (params?.knowledgeBaseId as string) ??
         (result.knowledgeBaseId as string) ??
         (data.knowledgeBaseId as string) ??
-        (params?.knowledgeBaseId as string)
+        (data.id as string)
       if (kbId) {
         const kbName =
           (data.name as string) ?? (result.knowledgeBaseName as string) ?? 'Knowledge Base'

apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.ts

@@ -28,6 +28,7 @@ import {
   deleteTagDefinition,
   getDocumentTagDefinitions,
   getNextAvailableSlot,
+  getTagDefinitionById,
   getTagUsageStats,
   updateTagDefinition,
 } from '@/lib/knowledge/tags/service'
@@ -525,6 +526,7 @@ export const knowledgeBaseServerTool: BaseServerTool<KnowledgeBaseArgs, Knowledg
             message: `Tag "${newTag.displayName}" created successfully`,
             data: {
               id: newTag.id,
+              knowledgeBaseId: args.knowledgeBaseId,
               tagSlot: newTag.tagSlot,
               displayName: newTag.displayName,
               fieldType: newTag.fieldType,
@@ -551,12 +553,21 @@ export const knowledgeBaseServerTool: BaseServerTool<KnowledgeBaseArgs, Knowledg
             }
           }
 
+          const existingTag = await getTagDefinitionById(args.tagDefinitionId)
+          if (!existingTag) {
+            return {
+              success: false,
+              message: `Tag definition with ID "${args.tagDefinitionId}" not found`,
+            }
+          }
+
           const requestId = crypto.randomUUID().slice(0, 8)
           assertNotAborted()
           const updatedTag = await updateTagDefinition(args.tagDefinitionId, updateData, requestId)
 
           logger.info('Tag definition updated via copilot', {
             tagId: args.tagDefinitionId,
+            knowledgeBaseId: existingTag.knowledgeBaseId,
             userId: context.userId,
           })
 
@@ -565,6 +576,7 @@ export const knowledgeBaseServerTool: BaseServerTool<KnowledgeBaseArgs, Knowledg
             message: `Tag "${updatedTag.displayName}" updated successfully`,
             data: {
               id: updatedTag.id,
+              knowledgeBaseId: existingTag.knowledgeBaseId,
               tagSlot: updatedTag.tagSlot,
               displayName: updatedTag.displayName,
               fieldType: updatedTag.fieldType,
@@ -605,6 +617,7 @@ export const knowledgeBaseServerTool: BaseServerTool<KnowledgeBaseArgs, Knowledg
             success: true,
             message: `Tag "${deleted.displayName}" deleted successfully. All document/chunk references cleared.`,
             data: {
+              knowledgeBaseId: args.knowledgeBaseId,
               tagSlot: deleted.tagSlot,
               displayName: deleted.displayName,
             },