fix(kb): canonicalize knowledge-base upload keys (#5096)

* fix(kb): canonicalize knowledge-base upload keys

* fix tests

Author: icecrasher321

apps/sim/app/api/files/presigned/route.test.ts

@@ -158,10 +158,10 @@ function setupFileApiMocks(
 
   storageServiceMockFns.mockHasCloudStorage.mockReturnValue(cloudEnabled)
   storageServiceMockFns.mockGeneratePresignedUploadUrl.mockImplementation(
-    async (opts: { fileName: string; context: string }) => {
+    async (opts: { fileName: string; context: string; customKey?: string }) => {
       const timestamp = Date.now()
       const safeFileName = opts.fileName.replace(/[^a-zA-Z0-9.-]/g, '_')
-      const key = `${opts.context}/${timestamp}-ik3a6w4-${safeFileName}`
+      const key = opts.customKey ?? `${opts.context}/${timestamp}-ik3a6w4-${safeFileName}`
       return {
         url: 'https://example.com/presigned-url',
         key,
@@ -369,7 +369,7 @@ describe('/api/files/presigned', () => {
       const data = await response.json()
 
       expect(response.status).toBe(200)
-      expect(data.fileInfo.key).toMatch(/^knowledge-base\/.*knowledge-doc\.pdf$/)
+      expect(data.fileInfo.key).toMatch(/^kb\/.*knowledge-doc\.pdf$/)
       expect(data.directUploadSupported).toBe(true)
     })
 

apps/sim/app/api/files/presigned/route.ts

@@ -9,6 +9,7 @@ import { CopilotFiles } from '@/lib/uploads'
 import type { StorageContext } from '@/lib/uploads/config'
 import { USE_BLOB_STORAGE } from '@/lib/uploads/config'
 import { generateExecutionFileKey } from '@/lib/uploads/contexts/execution/utils'
+import { generateKnowledgeBaseFileKey } from '@/lib/uploads/contexts/knowledge-base/knowledge-base-file-manager'
 import { generateWorkspaceFileKey } from '@/lib/uploads/contexts/workspace/workspace-file-manager'
 import { generatePresignedUploadUrl, hasCloudStorage } from '@/lib/uploads/core/storage-service'
 import { insertFileMetadata, recordKnowledgeBaseFileOwnership } from '@/lib/uploads/server/metadata'
@@ -268,12 +269,14 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         )
       }
 
+      const customKey = generateKnowledgeBaseFileKey(fileName)
       presignedUrlResponse = await generatePresignedUploadUrl({
         fileName,
         contentType,
         fileSize,
         context: 'knowledge-base',
         userId: sessionUserId,
+        customKey,
         expirationSeconds: 3600,
         metadata: { workspaceId },
       })

apps/sim/app/api/files/upload/route.ts

@@ -19,6 +19,7 @@ import {
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { captureServerEvent } from '@/lib/posthog/server'
 import type { StorageContext } from '@/lib/uploads/config'
+import { generateKnowledgeBaseFileKey } from '@/lib/uploads/contexts/knowledge-base/knowledge-base-file-manager'
 import { generateWorkspaceFileKey } from '@/lib/uploads/contexts/workspace/workspace-file-manager'
 import { MAX_WORKSPACE_FORMDATA_FILE_SIZE } from '@/lib/uploads/shared/types'
 import { isImageFileType, resolveFileType } from '@/lib/uploads/utils/file-utils'
@@ -155,9 +156,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
 
         logger.info(`Uploading knowledge-base file: ${originalName}`)
 
-        const timestamp = Date.now()
-        const safeFileName = sanitizeFileName(originalName)
-        const storageKey = `kb/${timestamp}-${safeFileName}`
+        const storageKey = generateKnowledgeBaseFileKey(originalName)
 
         const metadata: Record<string, string> = {
           originalName: originalName,

apps/sim/lib/uploads/contexts/knowledge-base/knowledge-base-file-manager.ts

@@ -0,0 +1,16 @@
+import { randomBytes } from 'crypto'
+import { sanitizeFileName } from '@/executor/constants'
+
+/**
+ * Generate a canonical knowledge-base storage key.
+ *
+ * Direct/presigned uploads previously used the generic `${context}/...` key
+ * shape (`knowledge-base/...`). New KB uploads should use the same `kb/...`
+ * prefix as server-side uploads so key-derived context inference is consistent.
+ */
+export function generateKnowledgeBaseFileKey(fileName: string): string {
+  const timestamp = Date.now()
+  const random = randomBytes(8).toString('hex')
+  const safeFileName = sanitizeFileName(fileName)
+  return `kb/${timestamp}-${random}-${safeFileName}`
+}

apps/sim/lib/uploads/utils/file-utils.test.ts

@@ -4,6 +4,7 @@
 import { createLogger } from '@sim/logger'
 import { describe, expect, it } from 'vitest'
 import {
+  inferContextFromKey,
   isAbortError,
   isInternalFileUrl,
   isNetworkError,
@@ -47,6 +48,32 @@ describe('isInternalFileUrl', () => {
   })
 })
 
+describe('inferContextFromKey', () => {
+  it('maps both kb/ and knowledge-base/ prefixes to knowledge-base', () => {
+    expect(inferContextFromKey('kb/1700000000000-doc.pdf')).toBe('knowledge-base')
+    // Direct/presigned uploads key as `${context}/...`, i.e. `knowledge-base/...`
+    expect(inferContextFromKey('knowledge-base/1781612506186-b2442e0dc045cb6c-doc.txt')).toBe(
+      'knowledge-base'
+    )
+  })
+
+  it('maps the remaining context prefixes', () => {
+    expect(inferContextFromKey('chat/x')).toBe('chat')
+    expect(inferContextFromKey('copilot/x')).toBe('copilot')
+    expect(inferContextFromKey('execution/ws/wf/ex/x')).toBe('execution')
+    expect(inferContextFromKey('workspace/ws/x')).toBe('workspace')
+    expect(inferContextFromKey('profile-pictures/x')).toBe('profile-pictures')
+    expect(inferContextFromKey('og-images/x')).toBe('og-images')
+    expect(inferContextFromKey('workspace-logos/x')).toBe('workspace-logos')
+    expect(inferContextFromKey('logs/x')).toBe('logs')
+  })
+
+  it('throws for empty or unrecognized keys', () => {
+    expect(() => inferContextFromKey('')).toThrow()
+    expect(() => inferContextFromKey('mystery/x')).toThrow()
+  })
+})
+
 describe('isAbortError', () => {
   it('returns true for AbortError-named errors', () => {
     const err = new Error('aborted')

apps/sim/lib/uploads/utils/file-utils.ts

@@ -565,15 +565,19 @@ export function isInternalFileUrl(fileUrl: string): boolean {
 }
 
 /**
- * Infer storage context from file key using explicit prefixes
- * All files must use prefixed keys
+ * Infer storage context from a file key using its prefix.
+ *
+ * All stored files use prefixed keys. Knowledge-base objects carry one of two
+ * prefixes: `kb/` (server-side uploads) or `knowledge-base/` (direct/presigned
+ * uploads, whose default key is `${context}/...`). Both map to the same
+ * `knowledge-base` context.
  */
 export function inferContextFromKey(key: string): StorageContext {
   if (!key) {
     throw new Error('Cannot infer context from empty key')
   }
 
-  if (key.startsWith('kb/')) return 'knowledge-base'
+  if (key.startsWith('kb/') || key.startsWith('knowledge-base/')) return 'knowledge-base'
   if (key.startsWith('chat/')) return 'chat'
   if (key.startsWith('copilot/')) return 'copilot'
   if (key.startsWith('execution/')) return 'execution'
@@ -584,7 +588,7 @@ export function inferContextFromKey(key: string): StorageContext {
   if (key.startsWith('logs/')) return 'logs'
 
   throw new Error(
-    `File key must start with a context prefix (kb/, chat/, copilot/, execution/, workspace/, profile-pictures/, og-images/, workspace-logos/, or logs/). Got: ${key}`
+    `File key must start with a context prefix (kb/, knowledge-base/, chat/, copilot/, execution/, workspace/, profile-pictures/, og-images/, workspace-logos/, or logs/). Got: ${key}`
   )
 }