diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts index 0c9b4be4b23..c426602e85c 100644 --- a/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts @@ -1,20 +1,15 @@ import { type NextRequest, NextResponse } from 'next/server' import { completeKnowledgeDocumentUploadContract } from '@/lib/api/contracts/knowledge/upload-sessions' import { parseRequest } from '@/lib/api/server' -import { generateRequestId } from '@/lib/core/utils/request' +import { PlatformEvents } from '@/lib/core/telemetry' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { completeUploadSession } from '@/lib/uploads/upload-session/service' -import { uploadSessionErrorResponse } from '@/app/api/files/uploads/utils' +import { completeKnowledgeDocumentUpload } from '@/lib/knowledge/application/upload-sessions' +import { captureServerEvent } from '@/lib/posthog/server' import { - requireKnowledgeDocumentUploadAccess, + knowledgeDocumentUploadErrorResponse, requireKnowledgeDocumentUploadActor, - resolveKnowledgeDocumentUploadAttribution, } from '@/app/api/knowledge/[id]/documents/uploads/utils' -import { - finalizeKnowledgeDocumentUpload, - getOwnedKnowledgeDocumentUpload, - toV2KnowledgeDocumentUpload, -} from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' +import { toV2KnowledgeDocumentUpload } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' interface KnowledgeDocumentUploadRouteParams { params: Promise<{ id: string; uploadId: string }> @@ -28,44 +23,46 @@ export const POST = withRouteHandler( if (!parsed.success) return parsed.response const { id: knowledgeBaseId, uploadId } = parsed.data.params const { workspaceId } = parsed.data.query - const access = await requireKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId: actor.id, - }) - if (access instanceof NextResponse) return access - const requestId = generateRequestId() try { - const upload = await getOwnedKnowledgeDocumentUpload({ - knowledgeBaseId, - uploadId, - workspaceId, - userId: actor.id, - uploadToken: parsed.data.headers['upload-token'], - }) - const completed = await completeUploadSession({ - session: upload, - finalize: (claimed) => - finalizeKnowledgeDocumentUpload({ - claimed, - knowledgeBaseId, - knowledgeBaseName: access.knowledgeBase.name, - workspaceId, - userId: actor.id, - resolveAttribution: () => - resolveKnowledgeDocumentUploadAttribution({ workspaceId, userId: actor.id }), - source: 'ui', - requestId, - request, - actorName: actor.name, - actorEmail: actor.email, - }), + const completed = await completeKnowledgeDocumentUpload.execute({ + principal: { kind: 'session', userId: actor.id, sessionId: actor.sessionId }, + input: { + knowledgeBaseId, + assertedWorkspaceId: workspaceId, + uploadId, + uploadToken: parsed.data.headers['upload-token'], + source: 'ui', + }, + request, }) + if (completed.value.created) { + captureServerEvent( + actor.id, + 'knowledge_base_document_uploaded', + { + knowledge_base_id: completed.knowledgeBaseId, + workspace_id: completed.workspaceId, + document_count: 1, + upload_type: 'single', + }, + { + groups: { workspace: completed.workspaceId }, + setOnce: { first_document_uploaded_at: new Date().toISOString() }, + } + ) + PlatformEvents.knowledgeBaseDocumentsUploaded({ + knowledgeBaseId: completed.knowledgeBaseId, + documentsCount: 1, + uploadType: 'single', + mimeType: completed.value.document.mimeType, + fileSize: completed.value.document.fileSize, + }) + } return NextResponse.json({ - data: toV2KnowledgeDocumentUpload(completed.session, completed.value), + data: toV2KnowledgeDocumentUpload(completed.session, completed.value.document), }) } catch (error) { - const classified = uploadSessionErrorResponse(error) + const classified = knowledgeDocumentUploadErrorResponse(error) if (classified) return classified throw error } diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts index da327ab4703..2ae1d2bc0e7 100644 --- a/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts @@ -2,13 +2,11 @@ import { type NextRequest, NextResponse } from 'next/server' import { createKnowledgeDocumentUploadPartUrlsContract } from '@/lib/api/contracts/knowledge/upload-sessions' import { parseRequest } from '@/lib/api/server' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { createUploadPartUrls } from '@/lib/uploads/upload-session/service' -import { uploadSessionErrorResponse } from '@/app/api/files/uploads/utils' +import { issueKnowledgeDocumentUploadParts } from '@/lib/knowledge/application/upload-sessions' import { - requireKnowledgeDocumentUploadAccess, + knowledgeDocumentUploadErrorResponse, requireKnowledgeDocumentUploadActor, } from '@/app/api/knowledge/[id]/documents/uploads/utils' -import { getOwnedKnowledgeDocumentUpload } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' interface KnowledgeDocumentUploadRouteParams { params: Promise<{ id: string; uploadId: string }> @@ -26,28 +24,21 @@ export const POST = withRouteHandler( if (!parsed.success) return parsed.response const { id: knowledgeBaseId, uploadId } = parsed.data.params const { workspaceId } = parsed.data.query - const access = await requireKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId: actor.id, - }) - if (access instanceof NextResponse) return access try { - const upload = await getOwnedKnowledgeDocumentUpload({ - knowledgeBaseId, - uploadId, - workspaceId, - userId: actor.id, - uploadToken: parsed.data.headers['upload-token'], - }) - const parts = await createUploadPartUrls({ - session: upload, - partNumbers: parsed.data.body.partNumbers, - localOrigin: request.nextUrl.origin, + const { parts } = await issueKnowledgeDocumentUploadParts.execute({ + principal: { kind: 'session', userId: actor.id, sessionId: actor.sessionId }, + input: { + knowledgeBaseId, + assertedWorkspaceId: workspaceId, + uploadId, + uploadToken: parsed.data.headers['upload-token'], + partNumbers: parsed.data.body.partNumbers, + }, + request, }) return NextResponse.json({ data: { parts } }) } catch (error) { - const classified = uploadSessionErrorResponse(error) + const classified = knowledgeDocumentUploadErrorResponse(error) if (classified) return classified throw error } diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/route.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/route.ts index 6a44d82d895..4f9f0d2c5b1 100644 --- a/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/route.ts +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/[uploadId]/route.ts @@ -2,16 +2,12 @@ import { type NextRequest, NextResponse } from 'next/server' import { abortKnowledgeDocumentUploadContract } from '@/lib/api/contracts/knowledge/upload-sessions' import { parseRequest } from '@/lib/api/server' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { uploadSessionErrorResponse } from '@/app/api/files/uploads/utils' +import { cancelKnowledgeDocumentUpload } from '@/lib/knowledge/application/upload-sessions' import { - requireKnowledgeDocumentUploadAccess, + knowledgeDocumentUploadErrorResponse, requireKnowledgeDocumentUploadActor, } from '@/app/api/knowledge/[id]/documents/uploads/utils' -import { - abortKnowledgeDocumentUpload, - getOwnedKnowledgeDocumentUpload, - toV2KnowledgeDocumentUpload, -} from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' +import { toV2KnowledgeDocumentUpload } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' interface KnowledgeDocumentUploadRouteParams { params: Promise<{ id: string; uploadId: string }> @@ -25,24 +21,20 @@ export const DELETE = withRouteHandler( if (!parsed.success) return parsed.response const { id: knowledgeBaseId, uploadId } = parsed.data.params const { workspaceId } = parsed.data.query - const access = await requireKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId: actor.id, - }) - if (access instanceof NextResponse) return access try { - const upload = await getOwnedKnowledgeDocumentUpload({ - knowledgeBaseId, - uploadId, - workspaceId, - userId: actor.id, - uploadToken: parsed.data.headers['upload-token'], + const aborted = await cancelKnowledgeDocumentUpload.execute({ + principal: { kind: 'session', userId: actor.id, sessionId: actor.sessionId }, + input: { + knowledgeBaseId, + assertedWorkspaceId: workspaceId, + uploadId, + uploadToken: parsed.data.headers['upload-token'], + }, + request, }) - const aborted = await abortKnowledgeDocumentUpload(upload, knowledgeBaseId) return NextResponse.json({ data: toV2KnowledgeDocumentUpload(aborted, null) }) } catch (error) { - const classified = uploadSessionErrorResponse(error) + const classified = knowledgeDocumentUploadErrorResponse(error) if (classified) return classified throw error } diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/control-routes.test.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/control-routes.test.ts new file mode 100644 index 00000000000..8a4b043cfff --- /dev/null +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/control-routes.test.ts @@ -0,0 +1,182 @@ +/** + * @vitest-environment node + */ +import { NextRequest } from 'next/server' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + cancel: vi.fn(), + captureServerEvent: vi.fn(), + complete: vi.fn(), + parts: vi.fn(), + platformEvent: vi.fn(), + requireActor: vi.fn(), +})) + +vi.mock('@/lib/knowledge/application/upload-sessions', () => ({ + cancelKnowledgeDocumentUpload: { execute: mocks.cancel }, + completeKnowledgeDocumentUpload: { execute: mocks.complete }, + issueKnowledgeDocumentUploadParts: { execute: mocks.parts }, +})) + +vi.mock('@/lib/core/telemetry', () => ({ + PlatformEvents: { knowledgeBaseDocumentsUploaded: mocks.platformEvent }, +})) +vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: mocks.captureServerEvent })) +vi.mock('@/app/api/knowledge/[id]/documents/uploads/utils', () => ({ + knowledgeDocumentUploadErrorResponse: vi.fn(() => null), + requireKnowledgeDocumentUploadActor: mocks.requireActor, +})) +vi.mock('@/app/api/v2/knowledge/[id]/documents/uploads/utils', () => ({ + toV2KnowledgeDocumentUpload: (_session: unknown, document: unknown) => ({ + id: 'upload-1', + knowledgeBaseId: 'kb-1', + status: document ? 'completed' : 'aborted', + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + expiresAt: '2026-08-05T00:00:00.000Z', + error: null, + document, + }), +})) + +import { POST as COMPLETE } from '@/app/api/knowledge/[id]/documents/uploads/[uploadId]/complete/route' +import { POST as PARTS } from '@/app/api/knowledge/[id]/documents/uploads/[uploadId]/parts/route' +import { DELETE as CANCEL } from '@/app/api/knowledge/[id]/documents/uploads/[uploadId]/route' + +const WORKSPACE_ID = '6fc7631d-88cd-46f8-9f0a-d4764daef7f8' +const PRINCIPAL = { kind: 'session', userId: 'user-1', sessionId: 'session-1' } +const SESSION = { id: 'upload-1', knowledgeBaseId: 'kb-1' } +const DOCUMENT = { + id: 'upload-1', + knowledgeBaseId: 'kb-1', + filename: 'guide.pdf', + fileSize: 1024, + mimeType: 'application/pdf', + chunkCount: 0, + tokenCount: 0, + characterCount: 0, + enabled: true, + uploadedAt: new Date('2026-08-03T21:01:00.000Z'), +} + +function routeContext() { + return { params: Promise.resolve({ id: 'kb-1', uploadId: 'upload-1' }) } +} + +function controlUrl(suffix = '') { + return `http://localhost:3000/api/knowledge/kb-1/documents/uploads/upload-1${suffix}?workspaceId=${WORKSPACE_ID}` +} + +describe('internal knowledge-document upload control routes', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.requireActor.mockResolvedValue({ id: 'user-1', sessionId: 'session-1' }) + mocks.parts.mockResolvedValue({ + parts: [ + { + partNumber: 1, + url: 'https://storage.example/1', + headers: {}, + expiresAt: '2026-08-04T21:00:00.000Z', + }, + ], + }) + mocks.cancel.mockResolvedValue(SESSION) + mocks.complete.mockResolvedValue({ + session: SESSION, + value: { document: DOCUMENT, created: true, knowledgeBaseName: 'Docs' }, + alreadyCompleted: false, + workspaceId: WORKSPACE_ID, + knowledgeBaseId: 'kb-1', + }) + }) + + it('delegates multipart part signing with the current session principal', async () => { + const request = new NextRequest(controlUrl('/parts'), { + method: 'POST', + headers: { 'content-type': 'application/json', 'upload-token': 'token' }, + body: JSON.stringify({ partNumbers: [1] }), + }) + + const response = await PARTS(request, routeContext()) + + expect(response.status).toBe(200) + expect(mocks.parts).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: WORKSPACE_ID, + uploadId: 'upload-1', + uploadToken: 'token', + partNumbers: [1], + }, + request, + }) + }) + + it('delegates cancellation with the current session principal', async () => { + const request = new NextRequest(controlUrl(), { + method: 'DELETE', + headers: { 'upload-token': 'token' }, + }) + + const response = await CANCEL(request, routeContext()) + + expect(response.status).toBe(200) + expect(mocks.cancel).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: WORKSPACE_ID, + uploadId: 'upload-1', + uploadToken: 'token', + }, + request, + }) + }) + + it('delegates completion and emits UI analytics only for a new document', async () => { + const request = new NextRequest(controlUrl('/complete'), { + method: 'POST', + headers: { 'upload-token': 'token' }, + }) + + const response = await COMPLETE(request, routeContext()) + + expect(response.status).toBe(200) + expect(mocks.complete).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: WORKSPACE_ID, + uploadId: 'upload-1', + uploadToken: 'token', + source: 'ui', + }, + request, + }) + expect(mocks.captureServerEvent).toHaveBeenCalledTimes(1) + expect(mocks.platformEvent).toHaveBeenCalledTimes(1) + }) + + it('does not duplicate UI analytics on an idempotent completion retry', async () => { + mocks.complete.mockResolvedValue({ + session: SESSION, + value: { document: DOCUMENT, created: false, knowledgeBaseName: 'Docs' }, + alreadyCompleted: true, + workspaceId: WORKSPACE_ID, + knowledgeBaseId: 'kb-1', + }) + const request = new NextRequest(controlUrl('/complete'), { + method: 'POST', + headers: { 'upload-token': 'token' }, + }) + + await COMPLETE(request, routeContext()) + + expect(mocks.captureServerEvent).not.toHaveBeenCalled() + expect(mocks.platformEvent).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/route.test.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/route.test.ts index ea79d0f4dc9..afc86b5d24a 100644 --- a/apps/sim/app/api/knowledge/[id]/documents/uploads/route.test.ts +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/route.test.ts @@ -1,35 +1,33 @@ /** * @vitest-environment node */ -import { NextRequest, NextResponse } from 'next/server' +import { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' -const { - mockCreateKnowledgeDocumentUploadSession, - mockRequireKnowledgeDocumentUploadAccess, - mockRequireKnowledgeDocumentUploadActor, - mockRequireKnowledgeDocumentUploadBilling, -} = vi.hoisted(() => ({ - mockCreateKnowledgeDocumentUploadSession: vi.fn(), - mockRequireKnowledgeDocumentUploadAccess: vi.fn(), - mockRequireKnowledgeDocumentUploadActor: vi.fn(), - mockRequireKnowledgeDocumentUploadBilling: vi.fn(), +const mocks = vi.hoisted(() => ({ + createUpload: vi.fn(), + requireActor: vi.fn(), +})) + +vi.mock('@/lib/knowledge/application/upload-sessions', () => ({ + createKnowledgeDocumentUpload: { execute: mocks.createUpload }, })) vi.mock('@/app/api/knowledge/[id]/documents/uploads/utils', () => ({ - requireKnowledgeDocumentUploadAccess: mockRequireKnowledgeDocumentUploadAccess, - requireKnowledgeDocumentUploadActor: mockRequireKnowledgeDocumentUploadActor, - requireKnowledgeDocumentUploadBilling: mockRequireKnowledgeDocumentUploadBilling, + knowledgeDocumentUploadErrorResponse: vi.fn(() => null), + requireKnowledgeDocumentUploadActor: mocks.requireActor, })) -vi.mock('@/app/api/files/uploads/utils', () => ({ uploadSessionErrorResponse: vi.fn() })) + vi.mock('@/app/api/v2/knowledge/[id]/documents/uploads/utils', () => ({ - createKnowledgeDocumentUploadSession: mockCreateKnowledgeDocumentUploadSession, toV2KnowledgeDocumentUpload: (session: Record) => ({ - ...session, + id: session.id, + knowledgeBaseId: session.knowledgeBaseId, + status: session.status, name: session.fileName, contentType: session.contentType, size: session.fileSize, expiresAt: '2026-08-05T00:00:00.000Z', + error: null, document: null, }), })) @@ -37,91 +35,70 @@ vi.mock('@/app/api/v2/knowledge/[id]/documents/uploads/utils', () => ({ import { POST } from '@/app/api/knowledge/[id]/documents/uploads/route' const WORKSPACE_ID = '6fc7631d-88cd-46f8-9f0a-d4764daef7f8' +const SESSION = { + id: 'upload-1', + knowledgeBaseId: 'kb-1', + status: 'uploading', + fileName: 'guide.pdf', + contentType: 'application/pdf', + fileSize: 1024, + uploadToken: 'token', + transfer: { + method: 'put' as const, + url: 'https://storage.example/upload', + headers: { 'content-type': 'application/pdf' }, + }, +} function request() { - return POST( - new NextRequest('http://localhost:3000/api/knowledge/kb-1/documents/uploads', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - workspaceId: WORKSPACE_ID, - name: 'guide.pdf', - contentType: 'application/pdf', - size: 1024, - tag1: 'product', - processingOptions: { recipe: 'default', lang: 'en' }, - }), + const request = new NextRequest('http://localhost:3000/api/knowledge/kb-1/documents/uploads', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + workspaceId: WORKSPACE_ID, + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + tag1: 'product', }), - { params: Promise.resolve({ id: 'kb-1' }) } - ) + }) + return { + request, + response: POST(request, { params: Promise.resolve({ id: 'kb-1' }) }), + } } describe('POST /api/knowledge/[id]/documents/uploads', () => { beforeEach(() => { vi.clearAllMocks() - mockRequireKnowledgeDocumentUploadActor.mockResolvedValue({ id: 'user-1' }) - mockRequireKnowledgeDocumentUploadAccess.mockResolvedValue({ - knowledgeBase: { id: 'kb-1', name: 'Docs', workspaceId: WORKSPACE_ID }, - }) - mockRequireKnowledgeDocumentUploadBilling.mockResolvedValue({ actorUserId: 'user-1' }) - mockCreateKnowledgeDocumentUploadSession.mockResolvedValue({ - id: 'upload-1', - knowledgeBaseId: 'kb-1', - status: 'uploading', - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - uploadToken: 'token', - error: null, - transfer: { - method: 'put', - url: 'https://storage.example/upload', - headers: { 'content-type': 'application/pdf' }, - }, + mocks.requireActor.mockResolvedValue({ + id: 'user-1', + sessionId: 'session-1', + name: 'User', + email: 'user@example.com', }) + mocks.createUpload.mockResolvedValue(SESSION) }) - it('authorizes and bills before allocating a first-party upload session', async () => { - const response = await request() + it('constructs a server-authored session principal and delegates creation', async () => { + const call = request() + const response = await call.response expect(response.status).toBe(201) - expect(mockRequireKnowledgeDocumentUploadAccess).toHaveBeenCalledWith({ - knowledgeBaseId: 'kb-1', - workspaceId: WORKSPACE_ID, - userId: 'user-1', - }) - expect(mockCreateKnowledgeDocumentUploadSession).toHaveBeenCalledWith({ - workspaceId: WORKSPACE_ID, - userId: 'user-1', - knowledgeBaseId: 'kb-1', - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - metadata: { - tag1: 'product', - processingOptions: { recipe: 'default', lang: 'en' }, + expect(mocks.createUpload).toHaveBeenCalledWith({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: WORKSPACE_ID, + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + metadata: { tag1: 'product' }, }, - localOrigin: 'http://localhost:3000', + request: call.request, }) - expect((await response.json()).data).toMatchObject({ - session: { id: 'upload-1', status: 'uploading', document: null }, - uploadToken: 'token', - transfer: { method: 'put', url: 'https://storage.example/upload' }, + expect(await response.json()).toMatchObject({ + data: { session: { id: 'upload-1' }, uploadToken: 'token' }, }) - expect(mockRequireKnowledgeDocumentUploadBilling.mock.invocationCallOrder[0]).toBeLessThan( - mockCreateKnowledgeDocumentUploadSession.mock.invocationCallOrder[0] - ) - }) - - it('does not bill or allocate storage when write access is denied', async () => { - mockRequireKnowledgeDocumentUploadAccess.mockResolvedValue( - NextResponse.json({ error: 'Forbidden' }, { status: 403 }) - ) - - const response = await request() - - expect(response.status).toBe(403) - expect(mockRequireKnowledgeDocumentUploadBilling).not.toHaveBeenCalled() - expect(mockCreateKnowledgeDocumentUploadSession).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/route.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/route.ts index 58a1c69c253..26a46296fe2 100644 --- a/apps/sim/app/api/knowledge/[id]/documents/uploads/route.ts +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/route.ts @@ -2,17 +2,12 @@ import { type NextRequest, NextResponse } from 'next/server' import { createKnowledgeDocumentUploadContract } from '@/lib/api/contracts/knowledge/upload-sessions' import { parseRequest } from '@/lib/api/server' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { validateFileType } from '@/lib/uploads/utils/validation' -import { uploadSessionErrorResponse } from '@/app/api/files/uploads/utils' +import { createKnowledgeDocumentUpload } from '@/lib/knowledge/application/upload-sessions' import { - requireKnowledgeDocumentUploadAccess, + knowledgeDocumentUploadErrorResponse, requireKnowledgeDocumentUploadActor, - requireKnowledgeDocumentUploadBilling, } from '@/app/api/knowledge/[id]/documents/uploads/utils' -import { - createKnowledgeDocumentUploadSession, - toV2KnowledgeDocumentUpload, -} from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' +import { toV2KnowledgeDocumentUpload } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' interface KnowledgeDocumentUploadsRouteParams { params: Promise<{ id: string }> @@ -26,31 +21,18 @@ export const POST = withRouteHandler( if (!parsed.success) return parsed.response const { id: knowledgeBaseId } = parsed.data.params const { workspaceId, name, contentType, size, ...metadata } = parsed.data.body - const access = await requireKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId: actor.id, - }) - if (access instanceof NextResponse) return access - const billing = await requireKnowledgeDocumentUploadBilling({ - workspaceId, - userId: actor.id, - }) - if (billing instanceof NextResponse) return billing - const fileTypeError = validateFileType(name, contentType) - if (fileTypeError) { - return NextResponse.json({ error: fileTypeError.message }, { status: 415 }) - } try { - const upload = await createKnowledgeDocumentUploadSession({ - workspaceId, - userId: actor.id, - knowledgeBaseId, - fileName: name, - contentType, - fileSize: size, - metadata, - localOrigin: request.nextUrl.origin, + const upload = await createKnowledgeDocumentUpload.execute({ + principal: { kind: 'session', userId: actor.id, sessionId: actor.sessionId }, + input: { + knowledgeBaseId, + assertedWorkspaceId: workspaceId, + name, + contentType, + size, + metadata, + }, + request, }) return NextResponse.json( { @@ -63,7 +45,7 @@ export const POST = withRouteHandler( { status: 201 } ) } catch (error) { - const classified = uploadSessionErrorResponse(error) + const classified = knowledgeDocumentUploadErrorResponse(error) if (classified) return classified throw error } diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/utils.test.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/utils.test.ts new file mode 100644 index 00000000000..80ae8df1f47 --- /dev/null +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/utils.test.ts @@ -0,0 +1,36 @@ +/** + * @vitest-environment node + */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ getSession: vi.fn() })) + +vi.mock('@/lib/auth', () => ({ getSession: mocks.getSession })) + +import { requireKnowledgeDocumentUploadActor } from '@/app/api/knowledge/[id]/documents/uploads/utils' + +describe('knowledge-document upload session authentication', () => { + beforeEach(() => vi.clearAllMocks()) + + it('returns the authoritative session id with the authenticated user', async () => { + mocks.getSession.mockResolvedValue({ + user: { id: 'user-1', name: 'User', email: 'user@example.com' }, + session: { id: 'session-1' }, + }) + + await expect(requireKnowledgeDocumentUploadActor()).resolves.toEqual({ + id: 'user-1', + sessionId: 'session-1', + name: 'User', + email: 'user@example.com', + }) + }) + + it('fails fast when authenticated state has no session id', async () => { + mocks.getSession.mockResolvedValue({ user: { id: 'user-1' }, session: {} }) + + await expect(requireKnowledgeDocumentUploadActor()).rejects.toThrow( + 'Authenticated session is missing its session ID' + ) + }) +}) diff --git a/apps/sim/app/api/knowledge/[id]/documents/uploads/utils.ts b/apps/sim/app/api/knowledge/[id]/documents/uploads/utils.ts index 450b17ecd0b..04aa96cc40a 100644 --- a/apps/sim/app/api/knowledge/[id]/documents/uploads/utils.ts +++ b/apps/sim/app/api/knowledge/[id]/documents/uploads/utils.ts @@ -1,15 +1,12 @@ import { NextResponse } from 'next/server' import { getSession } from '@/lib/auth' -import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution' -import { - checkAttributedUsageLimits, - resolveBillingAttribution, -} from '@/lib/billing/core/billing-attribution' -import type { KnowledgeBaseAccessResult } from '@/app/api/knowledge/utils' -import { checkKnowledgeBaseWriteAccess } from '@/app/api/knowledge/utils' +import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing' +import { KnowledgeDocumentUnsupportedMediaTypeError } from '@/lib/knowledge/application/upload-sessions' +import { uploadSessionErrorResponse } from '@/app/api/files/uploads/utils' export interface KnowledgeDocumentUploadActor { id: string + sessionId: string name?: string | null email?: string | null } @@ -21,53 +18,22 @@ export async function requireKnowledgeDocumentUploadActor(): Promise< if (!session?.user?.id) { return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }) } + const sessionId = session.session?.id + if (!sessionId) throw new Error('Authenticated session is missing its session ID') return { id: session.user.id, + sessionId, name: session.user.name, email: session.user.email, } } -export async function requireKnowledgeDocumentUploadAccess(params: { - knowledgeBaseId: string - workspaceId: string - userId: string -}): Promise<{ knowledgeBase: KnowledgeBaseAccessResult['knowledgeBase'] } | NextResponse> { - const access = await checkKnowledgeBaseWriteAccess(params.knowledgeBaseId, params.userId) - if (!access.hasAccess) { - return 'notFound' in access && access.notFound - ? NextResponse.json({ error: 'Knowledge base not found' }, { status: 404 }) - : NextResponse.json({ error: 'Forbidden' }, { status: 403 }) +export function knowledgeDocumentUploadErrorResponse(error: unknown): NextResponse | null { + if (error instanceof KnowledgeDocumentUnsupportedMediaTypeError) { + return NextResponse.json({ error: error.message }, { status: 415 }) } - if (access.knowledgeBase.workspaceId !== params.workspaceId) { - return NextResponse.json({ error: 'Knowledge base not found' }, { status: 404 }) + if (error instanceof KnowledgeUsageLimitExceededError) { + return NextResponse.json({ error: error.message }, { status: 402 }) } - return { knowledgeBase: access.knowledgeBase } -} - -export async function requireKnowledgeDocumentUploadBilling(params: { - workspaceId: string - userId: string -}): Promise { - const attribution = await resolveKnowledgeDocumentUploadAttribution(params) - const usage = await checkAttributedUsageLimits(attribution) - if (usage.isExceeded) { - return NextResponse.json( - { - error: usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.', - }, - { status: 402 } - ) - } - return attribution -} - -export function resolveKnowledgeDocumentUploadAttribution(params: { - workspaceId: string - userId: string -}): Promise { - return resolveBillingAttribution({ - actorUserId: params.userId, - workspaceId: params.workspaceId, - }) + return uploadSessionErrorResponse(error) } diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/[documentId]/route.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/[documentId]/route.ts index 079e4d18311..f12d79ddd31 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/[documentId]/route.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/[documentId]/route.ts @@ -1,119 +1,110 @@ -import { NextResponse } from 'next/server' import { - type V2KnowledgeDocument, v2DeleteKnowledgeDocumentContract, v2GetKnowledgeDocumentContract, } from '@/lib/api/contracts/v2/knowledge' -import { getKnowledgeDocument } from '@/lib/knowledge/documents/service' -import { performDeleteKnowledgeDocument } from '@/lib/knowledge/orchestration' -import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' -import { resolveKnowledgeBase, serializeDate } from '@/app/api/v1/knowledge/utils' -import type { RateLimitResult } from '@/app/api/v1/middleware' -import { v2Data, v2Error, v2ErrorForOrchestration } from '@/app/api/v2/lib/response' +import { + defineV2JsonRoute, + type V2ErrorPolicy, + v2ApiKeyAuth, + v2OrchestrationErrorPolicy, + v2RateLimits, +} from '@/lib/api/server/routes' +import { + deleteKnowledgeDocument, + readKnowledgeDocument, +} from '@/lib/knowledge/application/documents' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { captureServerEvent } from '@/lib/posthog/server' +import { serializeDate } from '@/app/api/v1/knowledge/utils' +import { v2Error } from '@/app/api/v2/lib/response' export const dynamic = 'force-dynamic' export const revalidate = 0 -/** - * Resolves a knowledge base via the shared v1 ownership invariant - * ({@link resolveKnowledgeBase}) and renders any failure in the v2 envelope. A - * `404` is always `NOT_FOUND`; a `403` is masked as `NOT_FOUND` on reads and - * surfaced as `FORBIDDEN` on writes. - */ -async function resolveKnowledgeBaseScoped( - id: string, - workspaceId: string, - userId: string, - rateLimit: RateLimitResult, - level: 'read' | 'write' -): Promise<{ kb: KnowledgeBaseWithCounts } | NextResponse> { - const result = await resolveKnowledgeBase(id, workspaceId, userId, rateLimit, level) - if (!(result instanceof NextResponse)) return result - if (result.status === 404) return v2Error('NOT_FOUND', 'Knowledge base not found') - return level === 'read' - ? v2Error('NOT_FOUND', 'Knowledge base not found') - : v2Error('FORBIDDEN', 'Access denied') +function toProcessingStatus(status: string): 'pending' | 'processing' | 'completed' | 'failed' { + switch (status) { + case 'pending': + case 'processing': + case 'completed': + case 'failed': + return status + default: + throw new Error(`Unexpected knowledge document processing status: ${status}`) + } } +const concealKnowledgeDocumentReadAuthorization = { + render(error) { + const response = v2OrchestrationErrorPolicy.render(error) + if (response?.status === 403) return v2Error('NOT_FOUND', 'Knowledge base not found') + return response + }, +} satisfies V2ErrorPolicy + /** GET /api/v2/knowledge/[id]/documents/[documentId] — Get document details. */ -export const GET = withPublicApiRouteHandler({ +export const GET = defineV2JsonRoute({ contract: v2GetKnowledgeDocumentContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ input, auth: { userId, rateLimit } }) => { - const { id: knowledgeBaseId, documentId } = input.params - - const result = await resolveKnowledgeBaseScoped( - knowledgeBaseId, - input.query.workspaceId, - userId, - rateLimit, - 'read' - ) - if (result instanceof NextResponse) return result - - const doc = await getKnowledgeDocument(knowledgeBaseId, documentId) - if (!doc) return v2Error('NOT_FOUND', 'Document not found') - - const documentDetail: V2KnowledgeDocument = { - id: doc.id, - knowledgeBaseId: doc.knowledgeBaseId, - filename: doc.filename, - fileSize: doc.fileSize, - mimeType: doc.mimeType, - processingStatus: doc.processingStatus as V2KnowledgeDocument['processingStatus'], - processingError: doc.processingError, - processingStartedAt: serializeDate(doc.processingStartedAt), - processingCompletedAt: serializeDate(doc.processingCompletedAt), - chunkCount: doc.chunkCount, - tokenCount: doc.tokenCount, - characterCount: doc.characterCount, - enabled: doc.enabled, - connectorId: doc.connectorId, - connectorType: doc.connectorType ?? null, - sourceUrl: doc.sourceUrl, - createdAt: serializeDate(doc.uploadedAt), - } - - return v2Data({ document: documentDetail }, { rateLimit }) - }, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.readDocument, + rateLimit: v2RateLimits.publicApi, + errorPolicy: concealKnowledgeDocumentReadAuthorization, + mapInput: ({ params, query }) => ({ + knowledgeBaseId: params.id, + documentId: params.documentId, + assertedWorkspaceId: query.workspaceId, + }), + useCase: readKnowledgeDocument, + present: ({ document }) => ({ + data: { + document: { + id: document.id, + knowledgeBaseId: document.knowledgeBaseId, + filename: document.filename, + fileSize: document.fileSize, + mimeType: document.mimeType, + processingStatus: toProcessingStatus(document.processingStatus), + processingError: document.processingError, + processingStartedAt: serializeDate(document.processingStartedAt), + processingCompletedAt: serializeDate(document.processingCompletedAt), + chunkCount: document.chunkCount, + tokenCount: document.tokenCount, + characterCount: document.characterCount, + enabled: document.enabled, + connectorId: document.connectorId, + connectorType: document.connectorType, + sourceUrl: document.sourceUrl, + createdAt: serializeDate(document.uploadedAt), + }, + }, + }), }) /** DELETE /api/v2/knowledge/[id]/documents/[documentId] — Delete a document. */ -export const DELETE = withPublicApiRouteHandler({ +export const DELETE = defineV2JsonRoute({ contract: v2DeleteKnowledgeDocumentContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ request, input, auth: { requestId, userId, rateLimit } }) => { - const { id: knowledgeBaseId, documentId } = input.params - - const result = await resolveKnowledgeBaseScoped( - knowledgeBaseId, - input.query.workspaceId, - userId, - rateLimit, - 'write' - ) - if (result instanceof NextResponse) return result - - const doc = await getKnowledgeDocument(knowledgeBaseId, documentId) - if (!doc) return v2Error('NOT_FOUND', 'Document not found') - - const outcome = await performDeleteKnowledgeDocument({ - knowledgeBase: { - id: knowledgeBaseId, - name: result.kb.name, - workspaceId: input.query.workspaceId, - }, - document: { id: documentId, filename: doc.filename }, - userId, - source: 'api', - requestId, - request, - }) - if (!outcome.success) { - return v2ErrorForOrchestration(outcome.errorCode, outcome.error) + auth: v2ApiKeyAuth, + operation: knowledgeOperations.deleteDocument, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + mapInput: ({ params, query }) => ({ + knowledgeBaseId: params.id, + documentId: params.documentId, + assertedWorkspaceId: query.workspaceId, + source: 'api', + }), + useCase: deleteKnowledgeDocument, + onSuccess: ({ principal, input }) => { + if (principal.kind === 'personal_api_key') { + captureServerEvent( + principal.userId, + 'knowledge_base_document_deleted', + { + knowledge_base_id: input.knowledgeBaseId, + workspace_id: input.assertedWorkspaceId ?? '', + }, + input.assertedWorkspaceId ? { groups: { workspace: input.assertedWorkspaceId } } : undefined + ) } - - return v2Data({ id: documentId, deleted: true as const }, { rateLimit }) }, + present: ({ id }) => ({ data: { id, deleted: true as const } }), }) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/route.test.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/route.test.ts new file mode 100644 index 00000000000..cbb5e0d88d2 --- /dev/null +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/route.test.ts @@ -0,0 +1,218 @@ +/** + * @vitest-environment node + */ +import { NextRequest } from 'next/server' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { + mockAuthenticate, + mockCheckPreAuth, + mockCheckRateLimit, + mockAdmitUpload, + mockUploadDocument, + mockReadFormData, + mockReadFile, + mockUploadWorkspaceFile, + mockPlatformUploaded, + mockCapture, +} = vi.hoisted(() => ({ + mockAuthenticate: vi.fn(), + mockCheckPreAuth: vi.fn(), + mockCheckRateLimit: vi.fn(), + mockAdmitUpload: vi.fn(), + mockUploadDocument: vi.fn(), + mockReadFormData: vi.fn(), + mockReadFile: vi.fn(), + mockUploadWorkspaceFile: vi.fn(), + mockPlatformUploaded: vi.fn(), + mockCapture: vi.fn(), +})) + +vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => ({ + authenticateV2ApiKey: mockAuthenticate, + V2ApiKeyUnauthenticatedError: class V2ApiKeyUnauthenticatedError extends Error {}, +})) + +vi.mock('@/lib/core/rate-limiter', () => ({ + getRateLimit: () => ({ maxTokens: 100, refillRate: 100, refillIntervalMs: 60_000 }), + RateLimiter: class RateLimiter { + checkRateLimitDirect(...args: unknown[]) { + return mockCheckPreAuth(...args) + } + + checkRateLimitDirectOrThrow(...args: unknown[]) { + return mockCheckRateLimit(...args) + } + }, +})) + +vi.mock('@/app/api/v2/lib/gate', () => ({ v2ApiGateError: vi.fn().mockResolvedValue(null) })) + +vi.mock('@/lib/knowledge/application/documents', () => ({ + listKnowledgeDocuments: { + operation: { id: 'knowledge.documents.list' }, + execute: vi.fn(), + }, + admitKnowledgeDocumentUpload: { + operation: { id: 'knowledge.documents.upload' }, + execute: mockAdmitUpload, + }, + uploadKnowledgeDocument: { + operation: { id: 'knowledge.documents.upload' }, + execute: mockUploadDocument, + }, +})) + +vi.mock('@/lib/core/utils/stream-limits', () => ({ + isPayloadSizeLimitError: () => false, + readFormDataWithLimit: mockReadFormData, + readFileToBufferWithLimit: mockReadFile, +})) + +vi.mock('@/lib/uploads/contexts/workspace', () => ({ + uploadWorkspaceFile: mockUploadWorkspaceFile, +})) + +vi.mock('@/lib/core/telemetry', () => ({ + PlatformEvents: { knowledgeBaseDocumentsUploaded: mockPlatformUploaded }, +})) + +vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: mockCapture })) + +import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing' +import { POST } from '@/app/api/v2/knowledge/[id]/documents/route' + +const WORKSPACE_ID = 'workspace-1' +const RATE_LIMIT_OK = { + allowed: true, + remaining: 99, + resetAt: new Date('2024-01-01T01:00:00Z'), + retryAfterMs: 0, +} +const PRINCIPAL = { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' } as const + +function buildRequest() { + return new NextRequest( + `http://localhost/api/v2/knowledge/kb-1/documents?workspaceId=${WORKSPACE_ID}`, + { method: 'POST', headers: { 'x-api-key': 'secret' }, body: 'multipart-placeholder' } + ) +} + +describe('POST /api/v2/knowledge/[id]/documents', () => { + beforeEach(() => { + vi.clearAllMocks() + mockCheckPreAuth.mockResolvedValue(RATE_LIMIT_OK) + mockCheckRateLimit.mockResolvedValue(RATE_LIMIT_OK) + mockAuthenticate.mockResolvedValue({ + principal: PRINCIPAL, + rolloutUserId: 'user-1', + rateLimitSubjectIds: ['api-key:key-1', 'user:user-1'], + rateLimitSubscription: null, + keyType: 'personal', + }) + mockAdmitUpload.mockResolvedValue({ + knowledgeBaseId: 'kb-1', + knowledgeBaseName: 'Support docs', + workspaceId: WORKSPACE_ID, + storageActorUserId: 'user-1', + }) + const formData = new FormData() + formData.set('file', new File(['hello'], 'support.txt', { type: 'text/plain' })) + mockReadFormData.mockResolvedValue(formData) + mockReadFile.mockResolvedValue(Buffer.from('hello')) + mockUploadWorkspaceFile.mockResolvedValue({ url: 's3://workspace/support.txt' }) + mockUploadDocument.mockResolvedValue({ + created: true, + document: { + id: 'doc-1', + knowledgeBaseId: 'kb-1', + filename: 'support.txt', + fileUrl: 's3://workspace/support.txt', + fileSize: 5, + mimeType: 'text/plain', + chunkCount: 0, + tokenCount: 0, + characterCount: 0, + enabled: true, + uploadedAt: new Date('2024-01-01T00:00:00Z'), + }, + }) + }) + + it('admits before buffering and reauthorizes durable registration with code-defined admission', async () => { + const request = buildRequest() + + const response = await POST(request, { params: Promise.resolve({ id: 'kb-1' }) }) + + expect(response.status).toBe(201) + expect(mockAdmitUpload.mock.invocationCallOrder[0]).toBeLessThan( + mockReadFormData.mock.invocationCallOrder[0] + ) + expect(mockAdmitUpload).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { knowledgeBaseId: 'kb-1', assertedWorkspaceId: WORKSPACE_ID }, + request, + }) + expect(mockUploadWorkspaceFile).toHaveBeenCalledWith( + WORKSPACE_ID, + 'user-1', + Buffer.from('hello'), + 'support.txt', + 'text/plain' + ) + expect(mockUploadDocument).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: WORKSPACE_ID, + document: { + filename: 'support.txt', + fileUrl: 's3://workspace/support.txt', + fileSize: 5, + mimeType: 'text/plain', + }, + startProcessing: true, + usageAdmission: 'pre_admitted', + source: 'api', + }, + request, + }) + expect(mockPlatformUploaded).toHaveBeenCalledOnce() + expect(mockCapture).toHaveBeenCalledWith( + 'user-1', + 'knowledge_base_document_uploaded', + expect.objectContaining({ knowledge_base_id: 'kb-1' }), + expect.any(Object) + ) + }) + + it('maps usage admission to the v2 error before multipart buffering', async () => { + mockAdmitUpload.mockRejectedValue(new KnowledgeUsageLimitExceededError('Upgrade required')) + + const response = await POST(buildRequest(), { params: Promise.resolve({ id: 'kb-1' }) }) + + expect(response.status).toBe(402) + expect(await response.json()).toEqual({ + error: { code: 'USAGE_LIMIT_EXCEEDED', message: 'Upgrade required' }, + }) + expect(mockReadFormData).not.toHaveBeenCalled() + expect(mockUploadWorkspaceFile).not.toHaveBeenCalled() + expect(mockUploadDocument).not.toHaveBeenCalled() + }) + + it('does not create human analytics for a workspace key', async () => { + mockAuthenticate.mockResolvedValue({ + principal: { kind: 'workspace_api_key', workspaceId: WORKSPACE_ID, keyId: 'key-2' }, + rolloutUserId: 'billing-owner', + rateLimitSubjectIds: ['api-key:key-2', `workspace:${WORKSPACE_ID}`], + rateLimitSubscription: null, + keyType: 'workspace', + }) + + const response = await POST(buildRequest(), { params: Promise.resolve({ id: 'kb-1' }) }) + + expect(response.status).toBe(201) + expect(mockPlatformUploaded).toHaveBeenCalledOnce() + expect(mockCapture).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/route.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/route.ts index c189de04389..5acd58ac0c5 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/route.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/route.ts @@ -1,37 +1,41 @@ +import type { NextRequest } from 'next/server' import { NextResponse } from 'next/server' import { type V2KnowledgeDocumentSummary, v2ListKnowledgeDocumentsContract, v2UploadKnowledgeDocumentContract, } from '@/lib/api/contracts/v2/knowledge' +import { parseRequest } from '@/lib/api/server' import { - checkAttributedUsageLimits, - resolveBillingAttribution, - resolveSystemBillingAttribution, -} from '@/lib/billing/core/billing-attribution' + defineV2JsonRoute, + type V2ErrorPolicy, + v2ApiKeyAuth, + v2OrchestrationErrorPolicy, + v2RateLimits, +} from '@/lib/api/server/routes' +import type { JsonRouteContext } from '@/lib/api/server/routes/types' +import { admitV2Request, V2RouteInfrastructureError } from '@/lib/api/server/routes/v2-json-route' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { PlatformEvents } from '@/lib/core/telemetry' import { isPayloadSizeLimitError, readFileToBufferWithLimit, readFormDataWithLimit, } from '@/lib/core/utils/stream-limits' -import { getDocuments } from '@/lib/knowledge/documents/service' -import type { DocumentSortField, SortOrder } from '@/lib/knowledge/documents/types' -import { performUploadKnowledgeDocument } from '@/lib/knowledge/orchestration' -import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types' +import { withRouteHandler } from '@/lib/core/utils/with-route-handler' +import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing' +import { + admitKnowledgeDocumentUpload, + listKnowledgeDocuments, + uploadKnowledgeDocument, +} from '@/lib/knowledge/application/documents' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { captureServerEvent } from '@/lib/posthog/server' import { uploadWorkspaceFile } from '@/lib/uploads/contexts/workspace' import { MAX_KNOWLEDGE_DOCUMENT_FILE_SIZE } from '@/lib/uploads/shared/types' import { validateFileType } from '@/lib/uploads/utils/validation' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' -import { resolveKnowledgeBase, serializeDate } from '@/app/api/v1/knowledge/utils' -import type { RateLimitResult } from '@/app/api/v1/middleware' -import { - decodeCursor, - encodeCursor, - v2CursorList, - v2Data, - v2Error, - v2ErrorForOrchestration, -} from '@/app/api/v2/lib/response' +import { serializeDate } from '@/app/api/v1/knowledge/utils' +import { decodeCursor, encodeCursor, v2Error, v2ValidationError } from '@/app/api/v2/lib/response' export const dynamic = 'force-dynamic' export const revalidate = 0 @@ -39,127 +43,109 @@ export const revalidate = 0 const MAX_FILE_SIZE = MAX_KNOWLEDGE_DOCUMENT_FILE_SIZE const MAX_MULTIPART_OVERHEAD_BYTES = 1024 * 1024 -/** - * Resolves a knowledge base via the shared v1 ownership invariant - * ({@link resolveKnowledgeBase}) and renders any failure in the v2 envelope. A - * `404` is always `NOT_FOUND`; a `403` is masked as `NOT_FOUND` on reads and - * surfaced as `FORBIDDEN` on writes. - */ -async function resolveKnowledgeBaseScoped( - id: string, - workspaceId: string, - userId: string, - rateLimit: RateLimitResult, - level: 'read' | 'write' -): Promise<{ kb: KnowledgeBaseWithCounts } | NextResponse> { - const result = await resolveKnowledgeBase(id, workspaceId, userId, rateLimit, level) - if (!(result instanceof NextResponse)) return result - if (result.status === 404) return v2Error('NOT_FOUND', 'Knowledge base not found') - return level === 'read' - ? v2Error('NOT_FOUND', 'Knowledge base not found') - : v2Error('FORBIDDEN', 'Access denied') +const concealKnowledgeDocumentListAuthorization = { + render(error) { + const response = v2OrchestrationErrorPolicy.render(error) + if (response?.status === 403) return v2Error('NOT_FOUND', 'Knowledge base not found') + return response + }, +} satisfies V2ErrorPolicy + +function toV2DocumentSummary(document: { + id: string + knowledgeBaseId: string + filename: string + fileSize: number + mimeType: string + processingStatus?: 'pending' | 'processing' | 'completed' | 'failed' + chunkCount: number + tokenCount: number + characterCount: number + enabled: boolean + uploadedAt: Date +}): V2KnowledgeDocumentSummary { + return { + id: document.id, + knowledgeBaseId: document.knowledgeBaseId, + filename: document.filename, + fileSize: document.fileSize, + mimeType: document.mimeType, + processingStatus: document.processingStatus ?? 'pending', + chunkCount: document.chunkCount, + tokenCount: document.tokenCount, + characterCount: document.characterCount, + enabled: document.enabled, + createdAt: serializeDate(document.uploadedAt), + } } /** GET /api/v2/knowledge/[id]/documents — List documents in a knowledge base. */ -export const GET = withPublicApiRouteHandler({ +export const GET = defineV2JsonRoute({ contract: v2ListKnowledgeDocumentsContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ input, auth: { requestId, userId, rateLimit } }) => { - const { workspaceId, limit, cursor, search, enabledFilter, sortBy, sortOrder } = input.query - const { id: knowledgeBaseId } = input.params - - const result = await resolveKnowledgeBaseScoped( - knowledgeBaseId, - workspaceId, - userId, - rateLimit, - 'read' - ) - if (result instanceof NextResponse) return result - - const decodedCursor = cursor ? decodeCursor<{ offset: number }>(cursor) : null + auth: v2ApiKeyAuth, + operation: knowledgeOperations.listDocuments, + rateLimit: v2RateLimits.publicApi, + errorPolicy: concealKnowledgeDocumentListAuthorization, + mapInput: ({ params, query }) => { + const decodedCursor = query.cursor ? decodeCursor<{ offset: number }>(query.cursor) : null if ( - cursor && + query.cursor && (!decodedCursor || !Number.isInteger(decodedCursor.offset) || decodedCursor.offset < 0) ) { - return v2Error('BAD_REQUEST', 'Invalid cursor') + throw new OrchestrationError('validation', 'Invalid cursor') + } + return { + knowledgeBaseId: params.id, + assertedWorkspaceId: query.workspaceId, + enabledFilter: query.enabledFilter, + search: query.search, + limit: query.limit, + offset: decodedCursor?.offset ?? 0, + sortBy: query.sortBy, + sortOrder: query.sortOrder, } - const offset = decodedCursor?.offset ?? 0 - - const documentsResult = await getDocuments( - knowledgeBaseId, - { - enabledFilter: enabledFilter === 'all' ? undefined : enabledFilter, - search, - limit, - offset, - sortBy: sortBy as DocumentSortField, - sortOrder: sortOrder as SortOrder, - }, - requestId - ) - - const documents: V2KnowledgeDocumentSummary[] = documentsResult.documents.map((doc) => ({ - id: doc.id, - knowledgeBaseId, - filename: doc.filename, - fileSize: doc.fileSize, - mimeType: doc.mimeType, - processingStatus: doc.processingStatus, - chunkCount: doc.chunkCount, - tokenCount: doc.tokenCount, - characterCount: doc.characterCount, - enabled: doc.enabled, - createdAt: serializeDate(doc.uploadedAt), - })) - - const nextCursor = documentsResult.pagination.hasMore - ? encodeCursor({ offset: offset + limit }) - : null - return v2CursorList(documents, nextCursor, { rateLimit }) }, + useCase: listKnowledgeDocuments, + present: ({ documents, pagination }) => ({ + data: documents.map(toV2DocumentSummary), + nextCursor: pagination.hasMore + ? encodeCursor({ offset: pagination.offset + pagination.limit }) + : null, + }), }) -/** - * POST /api/v2/knowledge/[id]/documents — Upload a document to a knowledge base. - * - * Authorization runs fully before the multipart body is buffered: the workspace - * is a contract-validated query param (not a form field as in v1), so an - * unauthorized caller never streams a file into memory. Order: rate limit → - * KB ownership (write) → usage gate → buffered multipart read. - */ -export const POST = withPublicApiRouteHandler({ - contract: v2UploadKnowledgeDocumentContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ request, input, auth: { requestId, userId, rateLimit } }) => { - try { - const { id: knowledgeBaseId } = input.params - const { workspaceId } = input.query - - const result = await resolveKnowledgeBaseScoped( - knowledgeBaseId, - workspaceId, - userId, - rateLimit, - 'write' +/** POST /api/v2/knowledge/[id]/documents — Upload a document to a knowledge base. */ +export const POST = withRouteHandler( + async (request: NextRequest, context) => { + if (request.method !== v2UploadKnowledgeDocumentContract.method) { + throw new Error( + `Route received ${request.method} for ${v2UploadKnowledgeDocumentContract.method} contract ${v2UploadKnowledgeDocumentContract.path}` ) - if (result instanceof NextResponse) return result - - /** - * Gate before storage and indexing. Workspace keys use the billed account - * and immutable payer from one read; personal keys preserve their human actor. - */ - const billingAttribution = - rateLimit.keyType === 'workspace' - ? await resolveSystemBillingAttribution(workspaceId) - : await resolveBillingAttribution({ actorUserId: userId, workspaceId }) - const usage = await checkAttributedUsageLimits(billingAttribution) - if (usage.isExceeded) { - return v2Error( - 'USAGE_LIMIT_EXCEEDED', - usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' - ) - } + } + + const routeAdmission = await admitV2Request( + request, + knowledgeOperations.uploadDocument, + v2ApiKeyAuth, + v2RateLimits.publicApi + ) + if (!routeAdmission.success) return routeAdmission.response + + const parsed = await parseRequest(v2UploadKnowledgeDocumentContract, request, context ?? {}, { + validationErrorResponse: v2ValidationError, + }) + if (!parsed.success) return parsed.response + + const { principal } = routeAdmission.auth + const { id: knowledgeBaseId } = parsed.data.params + const { workspaceId } = parsed.data.query + + try { + const uploadAdmission = await admitKnowledgeDocumentUpload.execute({ + principal, + input: { knowledgeBaseId, assertedWorkspaceId: workspaceId }, + request, + }) let formData: FormData try { @@ -176,9 +162,7 @@ export const POST = withPublicApiRouteHandler({ const rawFile = formData.get('file') const file = rawFile instanceof File ? rawFile : null - if (!file) { - return v2Error('BAD_REQUEST', 'file form field is required') - } + if (!file) return v2Error('BAD_REQUEST', 'file form field is required') if (file.size > MAX_FILE_SIZE) { return v2Error( @@ -197,57 +181,80 @@ export const POST = withPublicApiRouteHandler({ label: 'knowledge document file', }) const contentType = file.type || 'application/octet-stream' - const uploadedFile = await uploadWorkspaceFile( - workspaceId, - userId, + uploadAdmission.workspaceId, + uploadAdmission.storageActorUserId, buffer, file.name, contentType ) - const outcome = await performUploadKnowledgeDocument({ - knowledgeBase: { id: knowledgeBaseId, name: result.kb.name, workspaceId }, - document: { - filename: file.name, - fileUrl: uploadedFile.url, - fileSize: file.size, - mimeType: contentType, + const result = await uploadKnowledgeDocument.execute({ + principal, + input: { + knowledgeBaseId, + assertedWorkspaceId: workspaceId, + document: { + filename: file.name, + fileUrl: uploadedFile.url, + fileSize: file.size, + mimeType: contentType, + }, + startProcessing: true, + usageAdmission: 'pre_admitted', + source: 'api', }, - startProcessing: 'queue', - billingAttribution, - uploadedBy: billingAttribution.actorUserId, - userId, - source: 'api', - requestId, request, }) - if (!outcome.success) { - return v2ErrorForOrchestration(outcome.errorCode, outcome.error) - } - const newDocument = outcome.document - const document: V2KnowledgeDocumentSummary = { - id: newDocument.id, + PlatformEvents.knowledgeBaseDocumentsUploaded({ knowledgeBaseId, - filename: newDocument.filename, - fileSize: newDocument.fileSize, - mimeType: newDocument.mimeType, - processingStatus: 'pending', - chunkCount: 0, - tokenCount: 0, - characterCount: 0, - enabled: newDocument.enabled, - createdAt: serializeDate(newDocument.uploadedAt), + documentsCount: 1, + uploadType: 'single', + mimeType: contentType, + fileSize: file.size, + }) + if (principal.kind === 'personal_api_key') { + captureServerEvent( + principal.userId, + 'knowledge_base_document_uploaded', + { + knowledge_base_id: knowledgeBaseId, + workspace_id: workspaceId, + document_count: 1, + upload_type: 'single', + }, + { + groups: { workspace: workspaceId }, + setOnce: { first_document_uploaded_at: new Date().toISOString() }, + } + ) } - return v2Data({ document }, { rateLimit, status: 201 }) + const document = toV2DocumentSummary(result.document) + const body = v2UploadKnowledgeDocumentContract.response.schema.parse({ + data: { document }, + }) + return NextResponse.json(body, { + status: 201, + headers: { 'Cache-Control': 'private, no-store' }, + }) } catch (error) { + if (error instanceof KnowledgeUsageLimitExceededError) { + return v2Error('USAGE_LIMIT_EXCEEDED', error.message) + } if (isPayloadSizeLimitError(error)) { return v2Error('PAYLOAD_TOO_LARGE', error.message) } - + const response = v2OrchestrationErrorPolicy.render(error) + if (response) return response throw error } }, -}) + { + unhandledErrorResponse: ({ error }) => + error instanceof V2RouteInfrastructureError + ? v2Error('SERVICE_UNAVAILABLE', 'Service temporarily unavailable') + : v2Error('INTERNAL_ERROR', 'Internal server error'), + } +) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.test.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.test.ts index 62a81397f64..e1db73c90fe 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.test.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.test.ts @@ -4,85 +4,81 @@ import { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' -const { - mockCheckRateLimit, - mockCompleteUploadSession, - mockFinalizeKnowledgeDocumentUpload, - mockResolveKnowledgeDocumentUploadAccess, - mockResolveKnowledgeDocumentUploadAttribution, -} = vi.hoisted(() => ({ - mockCheckRateLimit: vi.fn(), - mockCompleteUploadSession: vi.fn(), - mockFinalizeKnowledgeDocumentUpload: vi.fn(), - mockResolveKnowledgeDocumentUploadAccess: vi.fn(), - mockResolveKnowledgeDocumentUploadAttribution: vi.fn(), +const mocks = vi.hoisted(() => ({ + authenticateV2ApiKey: vi.fn(), + captureServerEvent: vi.fn(), + checkRateLimitDirect: vi.fn(), + checkRateLimitDirectOrThrow: vi.fn(), + completeUpload: vi.fn(), + gate: vi.fn(), + platformEvent: vi.fn(), })) -vi.mock('@/app/api/v1/middleware', () => ({ checkRateLimit: mockCheckRateLimit })) -vi.mock('@/app/api/v2/lib/gate', () => ({ - v2ApiGateError: vi.fn().mockResolvedValue(null), +vi.mock('@/lib/knowledge/application/upload-sessions', () => ({ + completeKnowledgeDocumentUpload: { + operation: { + id: 'knowledge.documents.upload.complete', + minimumRole: 'write', + workspaceApiKey: 'allow', + }, + execute: mocks.completeUpload, + }, +})) + +vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => ({ + authenticateV2ApiKey: mocks.authenticateV2ApiKey, + V2ApiKeyUnauthenticatedError: class V2ApiKeyUnauthenticatedError extends Error {}, })) -vi.mock('@/lib/uploads/upload-session/service', () => ({ - completeUploadSession: mockCompleteUploadSession, + +vi.mock('@/lib/core/rate-limiter', () => ({ + getRateLimit: () => ({ maxTokens: 100, refillRate: 50, refillIntervalMs: 60_000 }), + RateLimiter: class RateLimiter { + checkRateLimitDirect = mocks.checkRateLimitDirect + checkRateLimitDirectOrThrow = mocks.checkRateLimitDirectOrThrow + }, +})) + +vi.mock('@/lib/core/telemetry', () => ({ + PlatformEvents: { knowledgeBaseDocumentsUploaded: mocks.platformEvent }, })) +vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: mocks.captureServerEvent })) +vi.mock('@/app/api/v2/lib/gate', () => ({ v2ApiGateError: mocks.gate })) vi.mock('@/app/api/v2/knowledge/[id]/documents/uploads/utils', () => ({ - finalizeKnowledgeDocumentUpload: mockFinalizeKnowledgeDocumentUpload, - getOwnedKnowledgeDocumentUpload: vi.fn(() => SESSION), - resolveKnowledgeDocumentUploadAccess: mockResolveKnowledgeDocumentUploadAccess, - resolveKnowledgeDocumentUploadAttribution: mockResolveKnowledgeDocumentUploadAttribution, - toV2KnowledgeDocumentUpload: (session: Record, document: unknown) => ({ - ...session, - name: session.fileName, - contentType: session.contentType, - size: session.fileSize, + toV2KnowledgeDocumentUpload: (_session: unknown, document: { id: string } | null) => ({ + id: 'upload-1', + knowledgeBaseId: 'kb-1', + status: 'completed', + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, expiresAt: '2026-08-04T21:00:00.000Z', - document, + error: null, + document: document + ? { + id: document.id, + knowledgeBaseId: 'kb-1', + filename: 'guide.pdf', + fileSize: 1024, + mimeType: 'application/pdf', + processingStatus: 'pending', + chunkCount: 0, + tokenCount: 0, + characterCount: 0, + enabled: true, + createdAt: '2026-08-03T21:01:00.000Z', + } + : null, }), + v2KnowledgeDocumentUploadError: vi.fn(() => null), })) -import { OrchestrationError } from '@/lib/core/orchestration/types' import { POST } from '@/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route' const WORKSPACE_ID = '6fc7631d-88cd-46f8-9f0a-d4764daef7f8' -const FILE_URL = '/api/files/serve/s3/kb%2Fguide.pdf?context=knowledge-base' -const SESSION = { - id: 'upload-1', - workspaceId: WORKSPACE_ID, - userId: 'user-1', - knowledgeBaseId: 'kb-1', - workflowId: null, - executionId: null, - purpose: 'knowledge_document', - method: 'multipart', - storageContext: 'knowledge-base', - storageKey: 'kb/guide.pdf', - finalKey: 'kb/guide.pdf', - storageProvider: 's3', - providerUploadId: 'provider-1', - providerObjectVersion: null, - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - partSize: 8 * 1024 * 1024, - partCount: 1, - status: 'uploading', - metadata: { - tag1: 'product', - processingOptions: { recipe: 'default', lang: 'en' }, - }, - uploadToken: 'token', - createdAt: new Date('2026-08-03T21:00:00.000Z'), - expiresAt: new Date('2026-08-04T21:00:00.000Z'), - completedFileId: null, - error: null, - completedAt: null, - updatedAt: new Date('2026-08-03T21:00:00.000Z'), -} as const const DOCUMENT = { id: 'upload-1', knowledgeBaseId: 'kb-1', filename: 'guide.pdf', - fileUrl: FILE_URL, fileSize: 1024, mimeType: 'application/pdf', chunkCount: 0, @@ -91,94 +87,106 @@ const DOCUMENT = { enabled: true, uploadedAt: new Date('2026-08-03T21:01:00.000Z'), } -const RATE_LIMIT = { - allowed: true, - userId: 'user-1', - keyType: 'workspace', - limit: 100, - remaining: 99, - resetAt: new Date('2026-08-03T22:00:00.000Z'), +const RESULT = { + session: { id: 'upload-1' }, + value: { document: DOCUMENT, created: true, knowledgeBaseName: 'Docs' }, + alreadyCompleted: false, + workspaceId: WORKSPACE_ID, + knowledgeBaseId: 'kb-1', +} + +function auth(principal: Record) { + return { + principal, + rolloutUserId: 'user-1', + rateLimitSubjectIds: ['api-key:key-1', 'user:user-1'] as const, + rateLimitSubscription: null, + keyType: + principal.kind === 'workspace_api_key' ? ('workspace' as const) : ('personal' as const), + } } function request() { - return POST( - new NextRequest( - `http://localhost:3000/api/v2/knowledge/kb-1/documents/uploads/upload-1/complete?workspaceId=${WORKSPACE_ID}`, - { - method: 'POST', - headers: { 'upload-token': 'token' }, - } - ), - { params: Promise.resolve({ id: 'kb-1', uploadId: 'upload-1' }) } + const request = new NextRequest( + `http://localhost:3000/api/v2/knowledge/kb-1/documents/uploads/upload-1/complete?workspaceId=${WORKSPACE_ID}`, + { method: 'POST', headers: { 'upload-token': 'token', 'x-api-key': 'secret' } } ) + return { + request, + response: POST(request, { + params: Promise.resolve({ id: 'kb-1', uploadId: 'upload-1' }), + }), + } } -describe('POST knowledge-document multipart completion', () => { +describe('POST knowledge-document upload completion', () => { beforeEach(() => { vi.clearAllMocks() - mockCheckRateLimit.mockResolvedValue(RATE_LIMIT) - mockResolveKnowledgeDocumentUploadAccess.mockResolvedValue({ - kb: { id: 'kb-1', name: 'Docs' }, - }) - mockResolveKnowledgeDocumentUploadAttribution.mockResolvedValue({ actorUserId: 'payer-1' }) - mockFinalizeKnowledgeDocumentUpload.mockResolvedValue({ - value: DOCUMENT, - completedFileId: DOCUMENT.id, + mocks.authenticateV2ApiKey.mockResolvedValue( + auth({ kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }) + ) + mocks.gate.mockResolvedValue(null) + mocks.checkRateLimitDirect.mockResolvedValue({ + allowed: true, + remaining: 599, + resetAt: new Date('2026-08-04T21:00:00.000Z'), }) - mockCompleteUploadSession.mockImplementation(async ({ session, finalize }) => { - const finalized = await finalize(session) - return { - session: { ...session, status: 'completed', completedFileId: finalized.completedFileId }, - value: finalized.value, - alreadyCompleted: false, - } + mocks.checkRateLimitDirectOrThrow.mockResolvedValue({ + allowed: true, + remaining: 99, + resetAt: new Date('2026-08-04T21:00:00.000Z'), }) + mocks.completeUpload.mockResolvedValue(RESULT) }) - it('delegates completion to the shared finalizer and returns the bound document', async () => { - const response = await request() + it('delegates completion and emits v2 analytics only for a newly created document', async () => { + const call = request() + const response = await call.response expect(response.status).toBe(200) - expect(await response.json()).toMatchObject({ data: { document: { id: 'upload-1' } } }) - expect(mockFinalizeKnowledgeDocumentUpload).toHaveBeenCalledWith( - expect.objectContaining({ - claimed: SESSION, + expect(mocks.completeUpload).toHaveBeenCalledWith({ + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + input: { knowledgeBaseId: 'kb-1', - knowledgeBaseName: 'Docs', - workspaceId: WORKSPACE_ID, - userId: 'user-1', + assertedWorkspaceId: WORKSPACE_ID, + uploadId: 'upload-1', + uploadToken: 'token', source: 'api', - }) - ) - expect(mockCompleteUploadSession).toHaveBeenCalledWith( - expect.objectContaining({ - session: SESSION, - }) + }, + request: call.request, + }) + expect(mocks.captureServerEvent).toHaveBeenCalledWith( + 'user-1', + 'knowledge_base_document_uploaded', + expect.objectContaining({ knowledge_base_id: 'kb-1', workspace_id: WORKSPACE_ID }), + expect.any(Object) ) + expect(mocks.platformEvent).toHaveBeenCalledTimes(1) + expect(await response.json()).toMatchObject({ data: { document: { id: 'upload-1' } } }) }) - it('resolves the payer lazily, only when the finalizer asks for one', async () => { - await request() - - expect(mockResolveKnowledgeDocumentUploadAttribution).not.toHaveBeenCalled() + it('does not duplicate analytics for an idempotent completion retry', async () => { + mocks.completeUpload.mockResolvedValue({ + ...RESULT, + value: { ...RESULT.value, created: false }, + alreadyCompleted: true, + }) - const { resolveAttribution } = mockFinalizeKnowledgeDocumentUpload.mock.calls[0][0] - await resolveAttribution() + const response = await request().response - expect(mockResolveKnowledgeDocumentUploadAttribution).toHaveBeenCalledWith({ - workspaceId: WORKSPACE_ID, - userId: 'user-1', - rateLimit: RATE_LIMIT, - }) + expect(response.status).toBe(200) + expect(mocks.captureServerEvent).not.toHaveBeenCalled() + expect(mocks.platformEvent).not.toHaveBeenCalled() }) - it('maps an orchestration failure from the finalizer onto its v2 status', async () => { - mockFinalizeKnowledgeDocumentUpload.mockRejectedValue( - new OrchestrationError('payload_too_large', 'Storage limit exceeded') + it('does not attribute a workspace-key event to the billing owner', async () => { + mocks.authenticateV2ApiKey.mockResolvedValue( + auth({ kind: 'workspace_api_key', workspaceId: WORKSPACE_ID, keyId: 'key-1' }) ) - const response = await request() + await request().response - expect(response.status).toBe(413) + expect(mocks.captureServerEvent).not.toHaveBeenCalled() + expect(mocks.platformEvent).toHaveBeenCalledTimes(1) }) }) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts index 05b111c7cbf..aa9bf3e6751 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/complete/route.ts @@ -1,61 +1,56 @@ -import { NextResponse } from 'next/server' import { v2CompleteKnowledgeDocumentUploadContract } from '@/lib/api/contracts/v2/knowledge' -import { completeUploadSession } from '@/lib/uploads/upload-session/service' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' +import { defineV2JsonRoute, v2ApiKeyAuth, v2RateLimits } from '@/lib/api/server/routes' +import { PlatformEvents } from '@/lib/core/telemetry' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { completeKnowledgeDocumentUpload } from '@/lib/knowledge/application/upload-sessions' +import { captureServerEvent } from '@/lib/posthog/server' import { - finalizeKnowledgeDocumentUpload, - getOwnedKnowledgeDocumentUpload, - resolveKnowledgeDocumentUploadAccess, - resolveKnowledgeDocumentUploadAttribution, toV2KnowledgeDocumentUpload, + v2KnowledgeDocumentUploadError, } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' -import { v2CaughtOrchestrationError, v2Data } from '@/app/api/v2/lib/response' -export const POST = withPublicApiRouteHandler({ +export const POST = defineV2JsonRoute({ contract: v2CompleteKnowledgeDocumentUploadContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ request, input, auth: { requestId, userId, rateLimit } }) => { - try { - const { id: knowledgeBaseId, uploadId } = input.params - const { workspaceId } = input.query - - const access = await resolveKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId, - rateLimit, - }) - if (access instanceof NextResponse) return access - - const session = await getOwnedKnowledgeDocumentUpload({ - knowledgeBaseId, - uploadId, - workspaceId, - userId, - uploadToken: input.headers['upload-token'], - }) - const result = await completeUploadSession({ - session, - finalize: (claimed) => - finalizeKnowledgeDocumentUpload({ - claimed, - knowledgeBaseId, - knowledgeBaseName: access.kb.name, - workspaceId, - userId, - resolveAttribution: () => - resolveKnowledgeDocumentUploadAttribution({ workspaceId, userId, rateLimit }), - source: 'api', - requestId, - request, - }), + auth: v2ApiKeyAuth, + operation: knowledgeOperations.uploadComplete, + rateLimit: v2RateLimits.publicApi, + errorPolicy: { render: v2KnowledgeDocumentUploadError }, + mapInput: ({ params, query, headers }) => ({ + knowledgeBaseId: params.id, + assertedWorkspaceId: query.workspaceId, + uploadId: params.uploadId, + uploadToken: headers['upload-token'], + source: 'api' as const, + }), + useCase: completeKnowledgeDocumentUpload, + onSuccess: ({ principal, result }) => { + if (result.value.created && principal.kind === 'personal_api_key') { + captureServerEvent( + principal.userId, + 'knowledge_base_document_uploaded', + { + knowledge_base_id: result.knowledgeBaseId, + workspace_id: result.workspaceId, + document_count: 1, + upload_type: 'single', + }, + { + groups: { workspace: result.workspaceId }, + setOnce: { first_document_uploaded_at: new Date().toISOString() }, + } + ) + } + if (result.value.created) { + PlatformEvents.knowledgeBaseDocumentsUploaded({ + knowledgeBaseId: result.knowledgeBaseId, + documentsCount: 1, + uploadType: 'single', + mimeType: result.value.document.mimeType, + fileSize: result.value.document.fileSize, }) - - return v2Data(toV2KnowledgeDocumentUpload(result.session, result.value), { rateLimit }) - } catch (error) { - const classified = v2CaughtOrchestrationError(error) - if (classified) return classified - throw error } }, + present: (result) => ({ + data: toV2KnowledgeDocumentUpload(result.session, result.value.document), + }), }) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts index 8b7320f397a..6640972b06b 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/parts/route.ts @@ -1,46 +1,22 @@ -import { NextResponse } from 'next/server' import { v2CreateKnowledgeDocumentUploadPartUrlsContract } from '@/lib/api/contracts/v2/knowledge' -import { createUploadPartUrls } from '@/lib/uploads/upload-session/service' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' -import { - getOwnedKnowledgeDocumentUpload, - resolveKnowledgeDocumentUploadAccess, -} from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' -import { v2CaughtOrchestrationError, v2Data } from '@/app/api/v2/lib/response' +import { defineV2JsonRoute, v2ApiKeyAuth, v2RateLimits } from '@/lib/api/server/routes' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { issueKnowledgeDocumentUploadParts } from '@/lib/knowledge/application/upload-sessions' +import { v2KnowledgeDocumentUploadError } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' -export const POST = withPublicApiRouteHandler({ +export const POST = defineV2JsonRoute({ contract: v2CreateKnowledgeDocumentUploadPartUrlsContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ request, input, auth: { userId, rateLimit } }) => { - try { - const { id: knowledgeBaseId, uploadId } = input.params - const { workspaceId } = input.query - - const access = await resolveKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId, - rateLimit, - }) - if (access instanceof NextResponse) return access - - const session = await getOwnedKnowledgeDocumentUpload({ - knowledgeBaseId, - uploadId, - workspaceId, - userId, - uploadToken: input.headers['upload-token'], - }) - const parts = await createUploadPartUrls({ - session, - partNumbers: input.body.partNumbers, - localOrigin: request.nextUrl.origin, - }) - return v2Data({ parts }, { rateLimit }) - } catch (error) { - const classified = v2CaughtOrchestrationError(error) - if (classified) return classified - throw error - } - }, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.uploadParts, + rateLimit: v2RateLimits.publicApi, + errorPolicy: { render: v2KnowledgeDocumentUploadError }, + mapInput: ({ params, query, headers, body }) => ({ + knowledgeBaseId: params.id, + assertedWorkspaceId: query.workspaceId, + uploadId: params.uploadId, + uploadToken: headers['upload-token'], + partNumbers: body.partNumbers, + }), + useCase: issueKnowledgeDocumentUploadParts, + present: ({ parts }) => ({ data: { parts } }), }) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/route.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/route.ts index a7ba8567f2a..87194ec5d55 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/route.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/route.ts @@ -1,43 +1,24 @@ -import { NextResponse } from 'next/server' import { v2AbortKnowledgeDocumentUploadContract } from '@/lib/api/contracts/v2/knowledge' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' +import { defineV2JsonRoute, v2ApiKeyAuth, v2RateLimits } from '@/lib/api/server/routes' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { cancelKnowledgeDocumentUpload } from '@/lib/knowledge/application/upload-sessions' import { - abortKnowledgeDocumentUpload, - getOwnedKnowledgeDocumentUpload, - resolveKnowledgeDocumentUploadAccess, toV2KnowledgeDocumentUpload, + v2KnowledgeDocumentUploadError, } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' -import { v2CaughtOrchestrationError, v2Data } from '@/app/api/v2/lib/response' -export const DELETE = withPublicApiRouteHandler({ +export const DELETE = defineV2JsonRoute({ contract: v2AbortKnowledgeDocumentUploadContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ input, auth: { userId, rateLimit } }) => { - try { - const { id: knowledgeBaseId, uploadId } = input.params - const { workspaceId } = input.query - - const access = await resolveKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId, - rateLimit, - }) - if (access instanceof NextResponse) return access - - const session = await getOwnedKnowledgeDocumentUpload({ - knowledgeBaseId, - uploadId, - workspaceId, - userId, - uploadToken: input.headers['upload-token'], - }) - const aborted = await abortKnowledgeDocumentUpload(session, knowledgeBaseId) - return v2Data(toV2KnowledgeDocumentUpload(aborted, null), { rateLimit }) - } catch (error) { - const classified = v2CaughtOrchestrationError(error) - if (classified) return classified - throw error - } - }, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.uploadCancel, + rateLimit: v2RateLimits.publicApi, + errorPolicy: { render: v2KnowledgeDocumentUploadError }, + mapInput: ({ params, query, headers }) => ({ + knowledgeBaseId: params.id, + assertedWorkspaceId: query.workspaceId, + uploadId: params.uploadId, + uploadToken: headers['upload-token'], + }), + useCase: cancelKnowledgeDocumentUpload, + present: (session) => ({ data: toV2KnowledgeDocumentUpload(session, null) }), }) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/control-routes.test.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/control-routes.test.ts new file mode 100644 index 00000000000..8f90f4b030b --- /dev/null +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/control-routes.test.ts @@ -0,0 +1,163 @@ +/** + * @vitest-environment node + */ +import { NextRequest } from 'next/server' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + authenticateV2ApiKey: vi.fn(), + cancel: vi.fn(), + checkRateLimitDirect: vi.fn(), + checkRateLimitDirectOrThrow: vi.fn(), + gate: vi.fn(), + parts: vi.fn(), +})) + +vi.mock('@/lib/knowledge/application/upload-sessions', () => ({ + cancelKnowledgeDocumentUpload: { + operation: { + id: 'knowledge.documents.upload.cancel', + minimumRole: 'write', + workspaceApiKey: 'allow', + }, + execute: mocks.cancel, + }, + issueKnowledgeDocumentUploadParts: { + operation: { + id: 'knowledge.documents.upload.parts', + minimumRole: 'write', + workspaceApiKey: 'allow', + }, + execute: mocks.parts, + }, +})) + +vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => ({ + authenticateV2ApiKey: mocks.authenticateV2ApiKey, + V2ApiKeyUnauthenticatedError: class V2ApiKeyUnauthenticatedError extends Error {}, +})) + +vi.mock('@/lib/core/rate-limiter', () => ({ + getRateLimit: () => ({ maxTokens: 100, refillRate: 50, refillIntervalMs: 60_000 }), + RateLimiter: class RateLimiter { + checkRateLimitDirect = mocks.checkRateLimitDirect + checkRateLimitDirectOrThrow = mocks.checkRateLimitDirectOrThrow + }, +})) + +vi.mock('@/app/api/v2/lib/gate', () => ({ v2ApiGateError: mocks.gate })) +vi.mock('@/app/api/v2/knowledge/[id]/documents/uploads/utils', () => ({ + toV2KnowledgeDocumentUpload: () => ({ + id: 'upload-1', + knowledgeBaseId: 'kb-1', + status: 'aborted', + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + expiresAt: '2026-08-04T21:00:00.000Z', + error: null, + document: null, + }), + v2KnowledgeDocumentUploadError: vi.fn(() => null), +})) + +import { POST as PARTS } from '@/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/parts/route' +import { DELETE as CANCEL } from '@/app/api/v2/knowledge/[id]/documents/uploads/[uploadId]/route' + +const WORKSPACE_ID = '6fc7631d-88cd-46f8-9f0a-d4764daef7f8' +const PRINCIPAL = { + kind: 'personal_api_key' as const, + userId: 'user-1', + keyId: 'key-1', +} + +function context() { + return { params: Promise.resolve({ id: 'kb-1', uploadId: 'upload-1' }) } +} + +function controlUrl(suffix = '') { + return `http://localhost:3000/api/v2/knowledge/kb-1/documents/uploads/upload-1${suffix}?workspaceId=${WORKSPACE_ID}` +} + +describe('v2 knowledge-document upload control routes', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.authenticateV2ApiKey.mockResolvedValue({ + principal: PRINCIPAL, + rolloutUserId: 'user-1', + rateLimitSubjectIds: ['api-key:key-1', 'user:user-1'], + rateLimitSubscription: null, + keyType: 'personal', + }) + mocks.gate.mockResolvedValue(null) + mocks.checkRateLimitDirect.mockResolvedValue({ + allowed: true, + remaining: 599, + resetAt: new Date('2026-08-04T21:00:00.000Z'), + }) + mocks.checkRateLimitDirectOrThrow.mockResolvedValue({ + allowed: true, + remaining: 99, + resetAt: new Date('2026-08-04T21:00:00.000Z'), + }) + mocks.parts.mockResolvedValue({ + parts: [ + { + partNumber: 1, + url: 'https://storage.example/1', + headers: {}, + expiresAt: '2026-08-04T21:00:00.000Z', + }, + ], + }) + mocks.cancel.mockResolvedValue({ id: 'upload-1' }) + }) + + it('delegates part signing with the authenticated API-key principal', async () => { + const request = new NextRequest(controlUrl('/parts'), { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'upload-token': 'token', + 'x-api-key': 'secret', + }, + body: JSON.stringify({ partNumbers: [1] }), + }) + + const response = await PARTS(request, context()) + + expect(response.status).toBe(200) + expect(mocks.parts).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: WORKSPACE_ID, + uploadId: 'upload-1', + uploadToken: 'token', + partNumbers: [1], + }, + request, + }) + }) + + it('delegates cancellation with the authenticated API-key principal', async () => { + const request = new NextRequest(controlUrl(), { + method: 'DELETE', + headers: { 'upload-token': 'token', 'x-api-key': 'secret' }, + }) + + const response = await CANCEL(request, context()) + + expect(response.status).toBe(200) + expect(mocks.cancel).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: WORKSPACE_ID, + uploadId: 'upload-1', + uploadToken: 'token', + }, + request, + }) + }) +}) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.test.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.test.ts index 0d0b0462332..652dd4404fb 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.test.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.test.ts @@ -1,138 +1,183 @@ /** * @vitest-environment node */ -import { NextRequest, NextResponse } from 'next/server' +import { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' -const { - mockCheckRateLimit, - mockCreateKnowledgeDocumentUploadSession, - mockResolveKnowledgeDocumentUploadAccess, - mockResolveKnowledgeDocumentUploadBilling, -} = vi.hoisted(() => ({ - mockCheckRateLimit: vi.fn(), - mockCreateKnowledgeDocumentUploadSession: vi.fn(), - mockResolveKnowledgeDocumentUploadAccess: vi.fn(), - mockResolveKnowledgeDocumentUploadBilling: vi.fn(), +const mocks = vi.hoisted(() => ({ + authenticateV2ApiKey: vi.fn(), + checkRateLimitDirect: vi.fn(), + checkRateLimitDirectOrThrow: vi.fn(), + createUpload: vi.fn(), + gate: vi.fn(), })) -vi.mock('@/app/api/v1/middleware', () => ({ checkRateLimit: mockCheckRateLimit })) -vi.mock('@/app/api/v2/lib/gate', () => ({ - v2ApiGateError: vi.fn().mockResolvedValue(null), +vi.mock('@/lib/knowledge/application/upload-sessions', () => ({ + createKnowledgeDocumentUpload: { + operation: { + id: 'knowledge.documents.upload.create', + minimumRole: 'write', + workspaceApiKey: 'allow', + }, + execute: mocks.createUpload, + }, })) + +vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => ({ + authenticateV2ApiKey: mocks.authenticateV2ApiKey, + V2ApiKeyUnauthenticatedError: class V2ApiKeyUnauthenticatedError extends Error {}, +})) + +vi.mock('@/lib/core/rate-limiter', () => ({ + getRateLimit: () => ({ maxTokens: 100, refillRate: 50, refillIntervalMs: 60_000 }), + RateLimiter: class RateLimiter { + checkRateLimitDirect = mocks.checkRateLimitDirect + checkRateLimitDirectOrThrow = mocks.checkRateLimitDirectOrThrow + }, +})) + +vi.mock('@/app/api/v2/lib/gate', () => ({ v2ApiGateError: mocks.gate })) + vi.mock('@/app/api/v2/knowledge/[id]/documents/uploads/utils', () => ({ - createKnowledgeDocumentUploadSession: mockCreateKnowledgeDocumentUploadSession, - resolveKnowledgeDocumentUploadAccess: mockResolveKnowledgeDocumentUploadAccess, - resolveKnowledgeDocumentUploadBilling: mockResolveKnowledgeDocumentUploadBilling, toV2KnowledgeDocumentUpload: (session: Record) => ({ - ...session, + id: session.id, + knowledgeBaseId: session.knowledgeBaseId, + status: session.status, name: session.fileName, contentType: session.contentType, size: session.fileSize, expiresAt: '2026-08-04T21:00:00.000Z', + error: null, document: null, }), + v2KnowledgeDocumentUploadError: vi.fn(() => null), })) import { POST } from '@/app/api/v2/knowledge/[id]/documents/uploads/route' const WORKSPACE_ID = '6fc7631d-88cd-46f8-9f0a-d4764daef7f8' -const RATE_LIMIT = { - allowed: true, - userId: 'user-1', - keyType: 'workspace', - limit: 100, - remaining: 99, - resetAt: new Date('2026-08-03T22:00:00.000Z'), +const PRINCIPAL = { + kind: 'workspace_api_key' as const, + workspaceId: WORKSPACE_ID, + keyId: 'key-1', +} +const AUTH = { + principal: PRINCIPAL, + rolloutUserId: 'billing-owner-1', + rateLimitSubjectIds: ['api-key:key-1', `workspace:${WORKSPACE_ID}`] as const, + rateLimitSubscription: null, + keyType: 'workspace' as const, +} +const SESSION = { + id: 'upload-1', + knowledgeBaseId: 'kb-1', + status: 'uploading', + fileName: 'guide.pdf', + contentType: 'application/pdf', + fileSize: 1024, + uploadToken: 'token', + transfer: { + method: 'put' as const, + url: 'https://storage.example/upload', + headers: { 'content-type': 'application/pdf' }, + }, } -function request() { - return POST( - new NextRequest('http://localhost:3000/api/v2/knowledge/kb-1/documents/uploads', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - workspaceId: WORKSPACE_ID, - name: 'guide.pdf', - contentType: 'application/pdf', - size: 1024, - tag1: 'product', - processingOptions: { recipe: 'default', lang: 'en' }, - }), - }), - { params: Promise.resolve({ id: 'kb-1' }) } - ) +function request(body: Record) { + const request = new NextRequest('http://localhost:3000/api/v2/knowledge/kb-1/documents/uploads', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'x-api-key': 'secret' }, + body: JSON.stringify(body), + }) + return { + request, + response: POST(request, { params: Promise.resolve({ id: 'kb-1' }) }), + } } describe('POST /api/v2/knowledge/[id]/documents/uploads', () => { beforeEach(() => { vi.clearAllMocks() - mockCheckRateLimit.mockResolvedValue(RATE_LIMIT) - mockResolveKnowledgeDocumentUploadAccess.mockResolvedValue({ - kb: { id: 'kb-1', name: 'Docs' }, + mocks.authenticateV2ApiKey.mockResolvedValue(AUTH) + mocks.gate.mockResolvedValue(null) + mocks.checkRateLimitDirect.mockResolvedValue({ + allowed: true, + remaining: 599, + resetAt: new Date('2026-08-04T21:00:00.000Z'), }) - mockResolveKnowledgeDocumentUploadBilling.mockResolvedValue({ actorUserId: 'user-1' }) - mockCreateKnowledgeDocumentUploadSession.mockResolvedValue({ - id: 'upload-1', - knowledgeBaseId: 'kb-1', - status: 'uploading', - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - uploadToken: 'token', - error: null, - transfer: { - method: 'put', - url: 'https://storage.example/upload', - headers: { 'content-type': 'application/pdf' }, - }, + mocks.checkRateLimitDirectOrThrow.mockResolvedValue({ + allowed: true, + remaining: 99, + resetAt: new Date('2026-08-04T21:00:00.000Z'), }) + mocks.createUpload.mockResolvedValue(SESSION) }) - it('authorizes the knowledge base and runs usage billing before accepting storage', async () => { - const response = await request() + it('delegates creation with the authenticated principal and asserted workspace', async () => { + const call = request({ + workspaceId: WORKSPACE_ID, + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + tag1: 'product', + processingOptions: { recipe: 'default', lang: 'en' }, + }) + const response = await call.response expect(response.status).toBe(201) - expect(mockResolveKnowledgeDocumentUploadAccess).toHaveBeenCalledWith( - expect.objectContaining({ + expect(mocks.createUpload).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { knowledgeBaseId: 'kb-1', - workspaceId: WORKSPACE_ID, - userId: 'user-1', - }) - ) - expect(mockResolveKnowledgeDocumentUploadBilling).toHaveBeenCalled() - expect(mockCreateKnowledgeDocumentUploadSession).toHaveBeenCalledWith({ - workspaceId: WORKSPACE_ID, - userId: 'user-1', - knowledgeBaseId: 'kb-1', - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - metadata: { - tag1: 'product', - processingOptions: { recipe: 'default', lang: 'en' }, + assertedWorkspaceId: WORKSPACE_ID, + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + metadata: { + tag1: 'product', + processingOptions: { recipe: 'default', lang: 'en' }, + }, }, - localOrigin: 'http://localhost:3000', + request: call.request, }) - expect((await response.json()).data).toMatchObject({ - session: { id: 'upload-1', status: 'uploading', document: null }, - uploadToken: 'token', - transfer: { method: 'put', url: 'https://storage.example/upload' }, + expect(await response.json()).toMatchObject({ + data: { + session: { id: 'upload-1', status: 'uploading', document: null }, + uploadToken: 'token', + }, }) - expect(mockResolveKnowledgeDocumentUploadBilling.mock.invocationCallOrder[0]).toBeLessThan( - mockCreateKnowledgeDocumentUploadSession.mock.invocationCallOrder[0] - ) }) - it('does not run billing or create provider state when knowledge write access is denied', async () => { - mockResolveKnowledgeDocumentUploadAccess.mockResolvedValue( - NextResponse.json({ error: { code: 'FORBIDDEN', message: 'Access denied' } }, { status: 403 }) - ) + it('authenticates and rate limits before parsing an invalid request', async () => { + const response = await request({ workspaceId: WORKSPACE_ID }).response - const response = await request() + expect(response.status).toBe(400) + expect(mocks.authenticateV2ApiKey).toHaveBeenCalledTimes(1) + expect(mocks.checkRateLimitDirectOrThrow).toHaveBeenCalledTimes(2) + expect(mocks.createUpload).not.toHaveBeenCalled() + }) + + it('rejects oversized or server-authored credential-binding body fields', async () => { + const oversized = await request({ + workspaceId: WORKSPACE_ID, + name: 'guide.pdf', + contentType: 'application/pdf', + size: 100 * 1024 * 1024 + 1, + }).response + const forgedBinding = await request({ + workspaceId: WORKSPACE_ID, + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + authBinding: { + version: 1, + workspaceId: WORKSPACE_ID, + principal: { kind: 'workspace_api_key', workspaceId: WORKSPACE_ID, keyId: 'forged' }, + }, + }).response - expect(response.status).toBe(403) - expect(mockResolveKnowledgeDocumentUploadBilling).not.toHaveBeenCalled() - expect(mockCreateKnowledgeDocumentUploadSession).not.toHaveBeenCalled() + expect(oversized.status).toBe(400) + expect(forgedBinding.status).toBe(400) + expect(mocks.createUpload).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.ts index 509a01f6de0..03f1ea7289d 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/route.ts @@ -1,65 +1,35 @@ -import { NextResponse } from 'next/server' import { v2CreateKnowledgeDocumentUploadContract } from '@/lib/api/contracts/v2/knowledge' -import { validateFileType } from '@/lib/uploads/utils/validation' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' +import { defineV2JsonRoute, v2ApiKeyAuth, v2RateLimits } from '@/lib/api/server/routes' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { createKnowledgeDocumentUpload } from '@/lib/knowledge/application/upload-sessions' import { - createKnowledgeDocumentUploadSession, - resolveKnowledgeDocumentUploadAccess, - resolveKnowledgeDocumentUploadBilling, toV2KnowledgeDocumentUpload, + v2KnowledgeDocumentUploadError, } from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' -import { v2CaughtOrchestrationError, v2Data, v2Error } from '@/app/api/v2/lib/response' -export const POST = withPublicApiRouteHandler({ +export const POST = defineV2JsonRoute({ contract: v2CreateKnowledgeDocumentUploadContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ request, input, auth: { userId, rateLimit } }) => { - try { - const { id: knowledgeBaseId } = input.params - const { workspaceId, name, contentType, size, ...metadata } = input.body - - const access = await resolveKnowledgeDocumentUploadAccess({ - knowledgeBaseId, - workspaceId, - userId, - rateLimit, - }) - if (access instanceof NextResponse) return access - - const billing = await resolveKnowledgeDocumentUploadBilling({ - workspaceId, - userId, - rateLimit, - }) - if (billing instanceof NextResponse) return billing - - const fileTypeError = validateFileType(name, contentType) - if (fileTypeError) { - return v2Error('UNSUPPORTED_MEDIA_TYPE', fileTypeError.message) - } - - const session = await createKnowledgeDocumentUploadSession({ - workspaceId, - userId, - knowledgeBaseId, - fileName: name, - contentType, - fileSize: size, - metadata, - localOrigin: request.nextUrl.origin, - }) - return v2Data( - { - session: toV2KnowledgeDocumentUpload(session, null), - uploadToken: session.uploadToken, - transfer: session.transfer, - }, - { rateLimit, status: 201 } - ) - } catch (error) { - const classified = v2CaughtOrchestrationError(error) - if (classified) return classified - throw error + auth: v2ApiKeyAuth, + operation: knowledgeOperations.uploadCreate, + rateLimit: v2RateLimits.publicApi, + errorPolicy: { render: v2KnowledgeDocumentUploadError }, + mapInput: ({ params, body }) => { + const { workspaceId, name, contentType, size, ...metadata } = body + return { + knowledgeBaseId: params.id, + assertedWorkspaceId: workspaceId, + name, + contentType, + size, + metadata, } }, + useCase: createKnowledgeDocumentUpload, + present: (session) => ({ + data: { + session: toV2KnowledgeDocumentUpload(session, null), + uploadToken: session.uploadToken, + transfer: session.transfer, + }, + }), }) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/utils.test.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/utils.test.ts deleted file mode 100644 index d7f25983576..00000000000 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/utils.test.ts +++ /dev/null @@ -1,254 +0,0 @@ -/** - * @vitest-environment node - */ -import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' -import type { UploadSessionRecord } from '@/lib/uploads/upload-session/service' - -const { - mockAbortUploadSession, - mockCreateUploadSession, - mockFindBoundKnowledgeDocument, - mockPerformUploadKnowledgeDocument, - mockRecordKnowledgeBaseFileOwnership, -} = vi.hoisted(() => ({ - mockAbortUploadSession: vi.fn(), - mockCreateUploadSession: vi.fn(), - mockFindBoundKnowledgeDocument: vi.fn(), - mockPerformUploadKnowledgeDocument: vi.fn(), - mockRecordKnowledgeBaseFileOwnership: vi.fn(), -})) - -vi.mock('@/lib/knowledge/orchestration', () => ({ - performUploadKnowledgeDocument: mockPerformUploadKnowledgeDocument, -})) -vi.mock('@/lib/knowledge/orchestration/documents', () => ({ - findBoundKnowledgeDocument: mockFindBoundKnowledgeDocument, -})) -vi.mock('@/lib/uploads/upload-session/service', () => ({ - abortUploadSession: mockAbortUploadSession, - createUploadSession: mockCreateUploadSession, - getOwnedUploadSession: vi.fn(), -})) -vi.mock('@/lib/uploads/server/metadata', () => ({ - recordKnowledgeBaseFileOwnership: mockRecordKnowledgeBaseFileOwnership, -})) - -import { - abortKnowledgeDocumentUpload, - createKnowledgeDocumentUploadSession, - finalizeKnowledgeDocumentUpload, - toV2KnowledgeDocumentUpload, -} from '@/app/api/v2/knowledge/[id]/documents/uploads/utils' - -const WORKSPACE_ID = '6fc7631d-88cd-46f8-9f0a-d4764daef7f8' -const CLAIMED: UploadSessionRecord = { - id: 'upload-1', - workspaceId: WORKSPACE_ID, - userId: 'user-1', - knowledgeBaseId: 'kb-1', - workflowId: null, - executionId: null, - purpose: 'knowledge_document', - method: 'multipart', - storageContext: 'knowledge-base', - storageKey: 'kb/guide.pdf', - finalKey: 'kb/guide.pdf', - storageProvider: 's3', - providerUploadId: 'provider-1', - providerObjectVersion: null, - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - partSize: 8 * 1024 * 1024, - partCount: 1, - status: 'uploading', - metadata: { tag1: 'product', processingOptions: { recipe: 'default', lang: 'en' } }, - uploadToken: 'token', - createdAt: new Date('2026-08-03T21:00:00.000Z'), - expiresAt: new Date('2026-08-04T21:00:00.000Z'), - completedFileId: null, - error: null, - completedAt: null, - updatedAt: new Date('2026-08-03T21:00:00.000Z'), -} -const DOCUMENT = { id: 'upload-1', knowledgeBaseId: 'kb-1', filename: 'guide.pdf' } - -function finalize(resolveAttribution = vi.fn().mockResolvedValue({ actorUserId: 'payer-1' })) { - return finalizeKnowledgeDocumentUpload({ - claimed: CLAIMED, - knowledgeBaseId: 'kb-1', - knowledgeBaseName: 'Docs', - workspaceId: WORKSPACE_ID, - userId: 'user-1', - resolveAttribution, - source: 'api', - requestId: 'req-1', - request: new NextRequest('http://localhost:3000/api/v2/knowledge/kb-1'), - }) -} - -function createSession() { - return createKnowledgeDocumentUploadSession({ - workspaceId: WORKSPACE_ID, - userId: 'user-1', - knowledgeBaseId: 'kb-1', - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - metadata: { tag1: 'product' }, - localOrigin: 'http://localhost:3000', - }) -} - -describe('createKnowledgeDocumentUploadSession', () => { - beforeEach(() => { - vi.clearAllMocks() - mockCreateUploadSession.mockResolvedValue(CLAIMED) - mockRecordKnowledgeBaseFileOwnership.mockResolvedValue(undefined) - mockAbortUploadSession.mockResolvedValue({ ...CLAIMED, status: 'aborted' }) - }) - - it('records the ownership binding before returning the upload token', async () => { - await expect(createSession()).resolves.toBe(CLAIMED) - - expect(mockCreateUploadSession).toHaveBeenCalledWith({ - workspaceId: WORKSPACE_ID, - userId: 'user-1', - knowledgeBaseId: 'kb-1', - purpose: 'knowledge_document', - fileName: 'guide.pdf', - contentType: 'application/pdf', - fileSize: 1024, - metadata: { tag1: 'product' }, - localOrigin: 'http://localhost:3000', - }) - expect(mockRecordKnowledgeBaseFileOwnership).toHaveBeenCalledWith({ - key: 'kb/guide.pdf', - userId: 'user-1', - workspaceId: WORKSPACE_ID, - originalName: 'guide.pdf', - contentType: 'application/pdf', - size: 1024, - }) - expect(mockCreateUploadSession.mock.invocationCallOrder[0]).toBeLessThan( - mockRecordKnowledgeBaseFileOwnership.mock.invocationCallOrder[0] - ) - }) - - it('aborts provider state when the ownership binding cannot be recorded', async () => { - mockRecordKnowledgeBaseFileOwnership.mockRejectedValue(new Error('database unavailable')) - - await expect(createSession()).rejects.toThrow('database unavailable') - expect(mockAbortUploadSession).toHaveBeenCalledWith(CLAIMED) - }) -}) - -describe('toV2KnowledgeDocumentUpload', () => { - it('does not expose reusable upload capabilities after session creation', () => { - const serialized = toV2KnowledgeDocumentUpload(CLAIMED, null) - - expect(serialized).not.toHaveProperty('uploadToken') - expect(serialized).not.toHaveProperty('partSize') - expect(serialized).not.toHaveProperty('partCount') - expect(serialized).not.toHaveProperty('transfer') - }) -}) - -describe('abortKnowledgeDocumentUpload', () => { - beforeEach(() => { - vi.clearAllMocks() - mockAbortUploadSession.mockResolvedValue({ ...CLAIMED, status: 'aborted' }) - }) - - it('aborts an upload that no document is bound to', async () => { - mockFindBoundKnowledgeDocument.mockResolvedValue({ status: 'absent' }) - - await expect(abortKnowledgeDocumentUpload(CLAIMED, 'kb-1')).resolves.toMatchObject({ - status: 'aborted', - }) - expect(mockAbortUploadSession).toHaveBeenCalledWith(CLAIMED) - }) - - it('refuses to abort once a document is bound, so committed bytes survive', async () => { - mockFindBoundKnowledgeDocument.mockResolvedValue({ status: 'bound', document: DOCUMENT }) - - await expect(abortKnowledgeDocumentUpload(CLAIMED, 'kb-1')).rejects.toThrow( - 'Upload has already been completed' - ) - expect(mockAbortUploadSession).not.toHaveBeenCalled() - }) -}) - -describe('finalizeKnowledgeDocumentUpload', () => { - beforeEach(() => { - vi.clearAllMocks() - mockFindBoundKnowledgeDocument.mockResolvedValue({ status: 'absent' }) - mockPerformUploadKnowledgeDocument.mockResolvedValue({ - success: true, - document: DOCUMENT, - created: true, - }) - }) - - it('creates the document, carrying session tags and processing options through', async () => { - const result = await finalize() - - expect(result).toEqual({ value: DOCUMENT, completedFileId: 'upload-1' }) - expect(mockPerformUploadKnowledgeDocument).toHaveBeenCalledWith( - expect.objectContaining({ - documentId: 'upload-1', - startProcessing: 'queue', - uploadedBy: 'payer-1', - processingOptions: { recipe: 'default', lang: 'en' }, - document: expect.objectContaining({ filename: 'guide.pdf', tag1: 'product' }), - }) - ) - }) - - it('answers a retry from the bound document without resolving a payer', async () => { - mockFindBoundKnowledgeDocument.mockResolvedValue({ status: 'bound', document: DOCUMENT }) - const resolveAttribution = vi.fn() - - const result = await finalize(resolveAttribution) - - expect(result).toEqual({ value: DOCUMENT, completedFileId: 'upload-1' }) - expect(resolveAttribution).not.toHaveBeenCalled() - expect(mockPerformUploadKnowledgeDocument).not.toHaveBeenCalled() - }) - - it('retains completed bytes for retry when document creation fails', async () => { - mockPerformUploadKnowledgeDocument.mockResolvedValue({ - success: false, - errorCode: 'payload_too_large', - error: 'Storage limit exceeded', - }) - - await expect(finalize()).rejects.toThrow('Storage limit exceeded') - expect(mockFindBoundKnowledgeDocument).toHaveBeenCalledTimes(1) - }) - - it('lets a retry converge when the first response fails after the document binds', async () => { - mockFindBoundKnowledgeDocument - .mockResolvedValueOnce({ status: 'absent' }) - .mockResolvedValueOnce({ status: 'bound', document: DOCUMENT }) - mockPerformUploadKnowledgeDocument.mockRejectedValue(new Error('audit sink exploded')) - - await expect(finalize()).rejects.toThrow('audit sink exploded') - await expect(finalize()).resolves.toEqual({ - value: DOCUMENT, - completedFileId: 'upload-1', - }) - expect(mockPerformUploadKnowledgeDocument).toHaveBeenCalledTimes(1) - }) - - it('rejects an upload id already bound to a different document without deleting anything', async () => { - mockFindBoundKnowledgeDocument.mockResolvedValue({ status: 'conflict' }) - const resolveAttribution = vi.fn() - - await expect(finalize(resolveAttribution)).rejects.toThrow( - 'Upload id is already bound to a different document' - ) - expect(resolveAttribution).not.toHaveBeenCalled() - }) -}) diff --git a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/utils.ts b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/utils.ts index 4ef4bf2d8fa..fbf8e030418 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/utils.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/documents/uploads/utils.ts @@ -1,136 +1,23 @@ -import type { NextRequest } from 'next/server' -import { NextResponse } from 'next/server' +import type { NextResponse } from 'next/server' import type { V2KnowledgeDocumentSummary, V2KnowledgeDocumentUpload, } from '@/lib/api/contracts/v2/knowledge' -import { v2KnowledgeDocumentUploadMetadataSchema } from '@/lib/api/contracts/v2/knowledge' -import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution' -import { - checkAttributedUsageLimits, - resolveBillingAttribution, - resolveSystemBillingAttribution, -} from '@/lib/billing/core/billing-attribution' -import { OrchestrationError } from '@/lib/core/orchestration/types' -import { performUploadKnowledgeDocument } from '@/lib/knowledge/orchestration' +import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing' +import { KnowledgeDocumentUnsupportedMediaTypeError } from '@/lib/knowledge/application/upload-sessions' import type { CreatedKnowledgeDocument } from '@/lib/knowledge/orchestration/documents' -import { findBoundKnowledgeDocument } from '@/lib/knowledge/orchestration/documents' -import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types' -import { recordKnowledgeBaseFileOwnership } from '@/lib/uploads/server/metadata' -import { - abortUploadSession, - type CreatedUploadSession, - createUploadSession, - getOwnedUploadSession, - type UploadSessionRecord, -} from '@/lib/uploads/upload-session/service' -import { resolveKnowledgeBase, serializeDate } from '@/app/api/v1/knowledge/utils' -import type { RateLimitResult } from '@/app/api/v1/middleware' -import { v2Error } from '@/app/api/v2/lib/response' +import type { UploadSessionRecord } from '@/lib/uploads/upload-session/service' +import { serializeDate } from '@/app/api/v1/knowledge/utils' +import { v2CaughtOrchestrationError, v2Error } from '@/app/api/v2/lib/response' -export async function resolveKnowledgeDocumentUploadAccess(params: { - knowledgeBaseId: string - workspaceId: string - userId: string - rateLimit: RateLimitResult -}): Promise<{ kb: KnowledgeBaseWithCounts } | NextResponse> { - const result = await resolveKnowledgeBase( - params.knowledgeBaseId, - params.workspaceId, - params.userId, - params.rateLimit, - 'write' - ) - if (!(result instanceof NextResponse)) return result - if (result.status === 404) return v2Error('NOT_FOUND', 'Knowledge base not found') - return v2Error('FORBIDDEN', 'Access denied') -} - -/** - * Resolves the payer for an upload without enforcing usage limits. Completion uses this - * because its bytes were already admitted when the session was created; re-running - * admission there would strand uploaded parts and fail idempotent completion retries. - */ -export async function resolveKnowledgeDocumentUploadAttribution(params: { - workspaceId: string - userId: string - rateLimit: RateLimitResult -}): Promise { - return params.rateLimit.keyType === 'workspace' - ? resolveSystemBillingAttribution(params.workspaceId) - : resolveBillingAttribution({ - actorUserId: params.userId, - workspaceId: params.workspaceId, - }) -} - -/** Admission check for a new upload session. Enforced only at session creation. */ -export async function resolveKnowledgeDocumentUploadBilling(params: { - workspaceId: string - userId: string - rateLimit: RateLimitResult -}): Promise { - const attribution = await resolveKnowledgeDocumentUploadAttribution(params) - const usage = await checkAttributedUsageLimits(attribution) - if (usage.isExceeded) { - return v2Error( - 'USAGE_LIMIT_EXCEEDED', - usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' - ) +export function v2KnowledgeDocumentUploadError(error: unknown): NextResponse | null { + if (error instanceof KnowledgeDocumentUnsupportedMediaTypeError) { + return v2Error('UNSUPPORTED_MEDIA_TYPE', error.message) } - return attribution -} - -export async function getOwnedKnowledgeDocumentUpload(params: { - knowledgeBaseId: string - uploadId: string - workspaceId: string - userId: string - uploadToken: string -}): Promise { - return getOwnedUploadSession({ - uploadId: params.uploadId, - workspaceId: params.workspaceId, - userId: params.userId, - purpose: 'knowledge_document', - knowledgeBaseId: params.knowledgeBaseId, - uploadToken: params.uploadToken, - }) -} - -/** - * Creates a knowledge-document upload and records its ownership binding before the token is - * returned. Failed or abandoned sessions can then be reclaimed by the knowledge-base orphan - * sweeper without racing a later document insert. - */ -export async function createKnowledgeDocumentUploadSession(params: { - workspaceId: string - userId: string - knowledgeBaseId: string - fileName: string - contentType: string - fileSize: number - metadata: Record - localOrigin: string -}): Promise { - const session = await createUploadSession({ - ...params, - purpose: 'knowledge_document', - }) - try { - await recordKnowledgeBaseFileOwnership({ - key: session.storageKey, - userId: params.userId, - workspaceId: params.workspaceId, - originalName: params.fileName, - contentType: params.contentType, - size: params.fileSize, - }) - } catch (error) { - await abortUploadSession(session) - throw error + if (error instanceof KnowledgeUsageLimitExceededError) { + return v2Error('USAGE_LIMIT_EXCEEDED', error.message) } - return session + return v2CaughtOrchestrationError(error) } export function toV2KnowledgeDocumentSummary( @@ -170,105 +57,3 @@ export function toV2KnowledgeDocumentUpload( document: document ? toV2KnowledgeDocumentSummary(document) : null, } } - -export function knowledgeDocumentFileUrl(session: UploadSessionRecord): string { - if (session.storageContext !== 'knowledge-base') { - throw new Error('Knowledge-document upload has an invalid storage context') - } - const providerPrefix = session.storageProvider === 'local' ? '' : `${session.storageProvider}/` - return `/api/files/serve/${providerPrefix}${encodeURIComponent(session.storageKey)}?context=knowledge-base` -} - -function knowledgeDocumentInputFor(session: UploadSessionRecord) { - const { processingOptions: _processingOptions, ...documentTags } = - v2KnowledgeDocumentUploadMetadataSchema.parse(session.metadata) - return { - filename: session.fileName, - fileUrl: knowledgeDocumentFileUrl(session), - fileSize: session.fileSize, - mimeType: session.contentType, - ...documentTags, - } -} - -/** - * Aborts an upload session, refusing once a document is bound to it. - * - * The document binding remains the domain-level completion authority while the upload row - * protects the provider object lifecycle. - */ -export async function abortKnowledgeDocumentUpload( - session: UploadSessionRecord, - knowledgeBaseId: string -): Promise { - const bound = await findBoundKnowledgeDocument({ - documentId: session.id, - knowledgeBaseId, - document: knowledgeDocumentInputFor(session), - }) - if (bound.status !== 'absent') { - throw new OrchestrationError('conflict', 'Upload has already been completed') - } - return abortUploadSession(session) -} - -/** - * Binds a completed upload session to its knowledge document. Shared by the public v2 - * and session-authenticated routes so both get identical completion semantics. - * - * Ordering is load-bearing. A retry is answered from the already-bound document before any - * work that can fail independently of the upload runs, so a payer that became unresolvable - * after the session was created cannot turn a valid retry into an error. The ownership binding - * is recorded before the upload token is issued, so failures retain retriable state and the - * delayed orphan sweeper reclaims sessions that never bind to a document. - */ -export async function finalizeKnowledgeDocumentUpload(params: { - claimed: UploadSessionRecord - knowledgeBaseId: string - knowledgeBaseName: string | null - workspaceId: string - userId: string - resolveAttribution: () => Promise - source: 'api' | 'ui' - requestId: string - request: NextRequest - actorName?: string | null - actorEmail?: string | null -}): Promise<{ value: CreatedKnowledgeDocument; completedFileId: string }> { - const { claimed, knowledgeBaseId, workspaceId, requestId } = params - const { processingOptions } = v2KnowledgeDocumentUploadMetadataSchema.parse(claimed.metadata) - const document = knowledgeDocumentInputFor(claimed) - - const bound = await findBoundKnowledgeDocument({ - documentId: claimed.id, - knowledgeBaseId, - document, - }) - if (bound.status === 'bound') { - return { value: bound.document, completedFileId: bound.document.id } - } - if (bound.status === 'conflict') { - throw new OrchestrationError('conflict', 'Upload id is already bound to a different document') - } - - const billingAttribution = await params.resolveAttribution() - const outcome = await performUploadKnowledgeDocument({ - knowledgeBase: { id: knowledgeBaseId, name: params.knowledgeBaseName, workspaceId }, - document, - documentId: claimed.id, - startProcessing: 'queue', - processingOptions, - billingAttribution, - uploadedBy: billingAttribution.actorUserId, - userId: params.userId, - ...(params.actorName ? { actorName: params.actorName } : {}), - ...(params.actorEmail ? { actorEmail: params.actorEmail } : {}), - source: params.source, - requestId, - request: params.request, - }) - if (!outcome.success) { - throw new OrchestrationError(outcome.errorCode, outcome.error) - } - return { value: outcome.document, completedFileId: outcome.document.id } -} diff --git a/apps/sim/app/api/v2/knowledge/[id]/route.ts b/apps/sim/app/api/v2/knowledge/[id]/route.ts index f2da8d298cb..a92d1b5fa85 100644 --- a/apps/sim/app/api/v2/knowledge/[id]/route.ts +++ b/apps/sim/app/api/v2/knowledge/[id]/route.ts @@ -1,151 +1,123 @@ -import { NextResponse } from 'next/server' import { v2DeleteKnowledgeBaseContract, v2GetKnowledgeBaseContract, v2UpdateKnowledgeBaseContract, } from '@/lib/api/contracts/v2/knowledge' -import { loadActiveFolderPathIndex } from '@/lib/folders/queries' import { - performDeleteKnowledgeBase, - performUpdateKnowledgeBase, -} from '@/lib/knowledge/orchestration' + defineV2JsonRoute, + type V2ErrorPolicy, + v2ApiKeyAuth, + v2OrchestrationErrorPolicy, + v2RateLimits, +} from '@/lib/api/server/routes' +import { PlatformEvents } from '@/lib/core/telemetry' +import { + deleteKnowledgeBaseOperation, + readKnowledgeBase, + updateKnowledgeBaseOperation, +} from '@/lib/knowledge/application/knowledge-bases' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' -import { formatKnowledgeBase, resolveKnowledgeBase } from '@/app/api/v1/knowledge/utils' -import type { RateLimitResult } from '@/app/api/v1/middleware' -import { folderPathForId, resolveFolderPathIdentity } from '@/app/api/v2/lib/folders' -import { v2Data, v2Error, v2ErrorForOrchestration } from '@/app/api/v2/lib/response' +import { v2Error } from '@/app/api/v2/lib/response' export const dynamic = 'force-dynamic' export const revalidate = 0 -/** - * Resolves a knowledge base via the shared v1 ownership invariant - * ({@link resolveKnowledgeBase}: workspace access + KB-belongs-to-workspace) and - * renders any failure in the v2 envelope. A `404` (missing KB or workspace - * mismatch) is always `NOT_FOUND`; a `403` (no workspace access) is masked as - * `NOT_FOUND` on reads so cross-workspace KB existence never leaks, and surfaced - * as `FORBIDDEN` on writes. - */ -async function resolveKnowledgeBaseScoped( - id: string, - workspaceId: string, - userId: string, - rateLimit: RateLimitResult, - level: 'read' | 'write' -): Promise<{ kb: KnowledgeBaseWithCounts } | NextResponse> { - const result = await resolveKnowledgeBase(id, workspaceId, userId, rateLimit, level) - if (!(result instanceof NextResponse)) return result - if (result.status === 404) return v2Error('NOT_FOUND', 'Knowledge base not found') - return level === 'read' - ? v2Error('NOT_FOUND', 'Knowledge base not found') - : v2Error('FORBIDDEN', 'Access denied') +function toV2KnowledgeBase(knowledgeBase: KnowledgeBaseWithCounts, folderPath: string) { + return { + id: knowledgeBase.id, + name: knowledgeBase.name, + description: knowledgeBase.description, + tokenCount: knowledgeBase.tokenCount, + embeddingModel: knowledgeBase.embeddingModel, + embeddingDimension: knowledgeBase.embeddingDimension, + chunkingConfig: { + maxSize: knowledgeBase.chunkingConfig.maxSize, + minSize: knowledgeBase.chunkingConfig.minSize, + overlap: knowledgeBase.chunkingConfig.overlap, + strategy: knowledgeBase.chunkingConfig.strategy, + strategyOptions: knowledgeBase.chunkingConfig.strategyOptions + ? { + pattern: knowledgeBase.chunkingConfig.strategyOptions.pattern, + separators: knowledgeBase.chunkingConfig.strategyOptions.separators, + recipe: knowledgeBase.chunkingConfig.strategyOptions.recipe, + strictBoundaries: knowledgeBase.chunkingConfig.strategyOptions.strictBoundaries, + } + : undefined, + }, + docCount: knowledgeBase.docCount, + connectorTypes: knowledgeBase.connectorTypes, + createdAt: knowledgeBase.createdAt.toISOString(), + updatedAt: knowledgeBase.updatedAt.toISOString(), + folderPath, + } } +const concealKnowledgeBaseReadAuthorization = { + render(error) { + const response = v2OrchestrationErrorPolicy.render(error) + if (response?.status === 403) return v2Error('NOT_FOUND', 'Knowledge base not found') + return response + }, +} satisfies V2ErrorPolicy + /** GET /api/v2/knowledge/[id] — Get knowledge base details. */ -export const GET = withPublicApiRouteHandler({ +export const GET = defineV2JsonRoute({ contract: v2GetKnowledgeBaseContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ input, auth: { userId, rateLimit } }) => { - const { id } = input.params - const result = await resolveKnowledgeBaseScoped( - id, - input.query.workspaceId, - userId, - rateLimit, - 'read' - ) - if (result instanceof NextResponse) return result - - const folderIndex = await loadActiveFolderPathIndex(input.query.workspaceId, 'knowledge_base') - - return v2Data( - { - knowledgeBase: { - ...formatKnowledgeBase(result.kb), - folderPath: folderPathForId(folderIndex, result.kb.folderId), - }, - }, - { rateLimit } - ) - }, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.read, + rateLimit: v2RateLimits.publicApi, + errorPolicy: concealKnowledgeBaseReadAuthorization, + mapInput: ({ params, query }) => ({ + knowledgeBaseId: params.id, + assertedWorkspaceId: query.workspaceId, + }), + useCase: readKnowledgeBase, + present: ({ knowledgeBase, folderPath }) => ({ + data: { knowledgeBase: toV2KnowledgeBase(knowledgeBase, folderPath) }, + }), }) /** PUT /api/v2/knowledge/[id] — Update a knowledge base. */ -export const PUT = withPublicApiRouteHandler({ +export const PUT = defineV2JsonRoute({ contract: v2UpdateKnowledgeBaseContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ request, input, auth: { requestId, userId, rateLimit } }) => { - const { id } = input.params - const { workspaceId, name, description, chunkingConfig, folderPath } = input.body - - const result = await resolveKnowledgeBaseScoped(id, workspaceId, userId, rateLimit, 'write') - if (result instanceof NextResponse) return result - - const resolution = - folderPath === undefined - ? undefined - : await resolveFolderPathIdentity({ - workspaceId, - resourceType: 'knowledge_base', - path: folderPath, - }) - if (resolution && !resolution.found) { - return v2Error('NOT_FOUND', 'Folder not found') - } - - const outcome = await performUpdateKnowledgeBase({ - knowledgeBaseId: id, - workspaceId, - userId, - source: 'api', - updates: { name, description, chunkingConfig, folderId: resolution?.folderId }, - requestId, - request, - }) - if (!outcome.success) { - return v2ErrorForOrchestration(outcome.errorCode, outcome.error) - } - - const folderIndex = await loadActiveFolderPathIndex(workspaceId, 'knowledge_base') - return v2Data( - { - knowledgeBase: { - ...formatKnowledgeBase(outcome.knowledgeBase), - folderPath: folderPathForId(folderIndex, outcome.knowledgeBase.folderId), - }, - }, - { rateLimit } - ) + auth: v2ApiKeyAuth, + operation: knowledgeOperations.update, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + parseOptions: { + invalidJsonResponse: () => v2Error('BAD_REQUEST', 'Request body must be valid JSON'), }, + mapInput: ({ params, body }) => ({ + knowledgeBaseId: params.id, + assertedWorkspaceId: body.workspaceId, + name: body.name, + description: body.description, + chunkingConfig: body.chunkingConfig, + folderPath: body.folderPath, + source: 'api', + }), + useCase: updateKnowledgeBaseOperation, + present: ({ knowledgeBase, folderPath }) => ({ + data: { knowledgeBase: toV2KnowledgeBase(knowledgeBase, folderPath) }, + }), }) /** DELETE /api/v2/knowledge/[id] — Delete a knowledge base. */ -export const DELETE = withPublicApiRouteHandler({ +export const DELETE = defineV2JsonRoute({ contract: v2DeleteKnowledgeBaseContract, - rateLimitEndpoint: 'knowledge-detail', - handler: async ({ request, input, auth: { requestId, userId, rateLimit } }) => { - const { id } = input.params - const result = await resolveKnowledgeBaseScoped( - id, - input.query.workspaceId, - userId, - rateLimit, - 'write' - ) - if (result instanceof NextResponse) return result - - const outcome = await performDeleteKnowledgeBase({ - knowledgeBase: { id, name: result.kb.name, workspaceId: input.query.workspaceId }, - userId, - source: 'api', - requestId, - request, - }) - if (!outcome.success) { - return v2ErrorForOrchestration(outcome.errorCode, outcome.error) - } - - return v2Data({ id, deleted: true as const }, { rateLimit }) + auth: v2ApiKeyAuth, + operation: knowledgeOperations.delete, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + mapInput: ({ params, query }) => ({ + knowledgeBaseId: params.id, + assertedWorkspaceId: query.workspaceId, + source: 'api', + }), + useCase: deleteKnowledgeBaseOperation, + onSuccess: ({ result }) => { + PlatformEvents.knowledgeBaseDeleted({ knowledgeBaseId: result.id }) }, + present: ({ id }) => ({ data: { id, deleted: true as const } }), }) diff --git a/apps/sim/app/api/v2/knowledge/folders/route.ts b/apps/sim/app/api/v2/knowledge/folders/route.ts index 3d5ff54a79a..aac1e97d802 100644 --- a/apps/sim/app/api/v2/knowledge/folders/route.ts +++ b/apps/sim/app/api/v2/knowledge/folders/route.ts @@ -5,123 +5,98 @@ import { v2RelocateKnowledgeFolderContract, } from '@/lib/api/contracts/v2/knowledge' import { - createFolderAtPath, - deleteFolderByPath, - relocateFolderByPath, -} from '@/lib/folders/orchestration' -import { listActiveFolderRows, loadActiveFolderPathIndex } from '@/lib/folders/queries' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' -import { resolveWorkspaceAccess } from '@/app/api/v1/middleware' + defineV2JsonRoute, + v2ApiKeyAuth, + v2OrchestrationErrorPolicy, + v2RateLimits, +} from '@/lib/api/server/routes' +import { toFolderPathView } from '@/lib/folders/paths' import { - resolveFolderPathId, - toV2PathFolder, - v2FolderPathMutationError, -} from '@/app/api/v2/lib/folders' -import { v2CursorList, v2Data, v2Error, v2WorkspaceAccessError } from '@/app/api/v2/lib/response' + createKnowledgeFolder, + deleteKnowledgeFolder, + listKnowledgeFolders, + relocateKnowledgeFolder, +} from '@/lib/knowledge/application/folders' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { v2Error } from '@/app/api/v2/lib/response' export const dynamic = 'force-dynamic' export const revalidate = 0 -export const GET = withPublicApiRouteHandler({ +export const GET = defineV2JsonRoute({ contract: v2ListKnowledgeFoldersContract, - rateLimitEndpoint: 'knowledge', - handler: async ({ input, auth: { userId, rateLimit } }) => { - const { workspaceId, parentPath, search, sortBy, sortOrder } = input.query - const access = await resolveWorkspaceAccess(rateLimit, userId, workspaceId, 'read') - if (access) return v2WorkspaceAccessError(access) - - const index = await loadActiveFolderPathIndex(workspaceId, 'knowledge_base') - const parentId = parentPath === undefined ? undefined : resolveFolderPathId(index, parentPath) - if (parentPath !== undefined && parentId === undefined) { - return v2Error('NOT_FOUND', 'Folder not found') - } - const rows = await listActiveFolderRows(workspaceId, 'knowledge_base', { - parentId, - search, - sortBy, - sortOrder, - }) - return v2CursorList( - rows.map((row) => toV2PathFolder(row, index, false)), - null, - { rateLimit } - ) - }, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.listFolders, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + mapInput: ({ query }) => ({ + workspaceId: query.workspaceId, + parentPath: query.parentPath, + search: query.search, + sortBy: query.sortBy, + sortOrder: query.sortOrder, + }), + useCase: listKnowledgeFolders, + present: ({ folders }) => ({ + data: folders.map((folder) => toFolderPathView(folder, folder.path)), + nextCursor: null, + }), }) -export const POST = withPublicApiRouteHandler({ +export const POST = defineV2JsonRoute({ contract: v2CreateKnowledgeFolderContract, - rateLimitEndpoint: 'knowledge', - handler: async ({ input, auth: { userId, rateLimit } }) => { - const { workspaceId, path } = input.body - const access = await resolveWorkspaceAccess(rateLimit, userId, workspaceId, 'write') - if (access) return v2WorkspaceAccessError(access) - const result = await createFolderAtPath({ - resourceType: 'knowledge_base', - workspaceId, - userId, - path, - }) - if (!result.success || !result.folder) { - return v2FolderPathMutationError(result.errorCode, result.error ?? 'Failed to create folder') - } - const index = await loadActiveFolderPathIndex(workspaceId, 'knowledge_base') - return v2Data( - { folder: toV2PathFolder(result.folder, index, false) }, - { rateLimit, status: 201 } - ) + auth: v2ApiKeyAuth, + operation: knowledgeOperations.createFolder, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + parseOptions: { + invalidJsonResponse: () => v2Error('BAD_REQUEST', 'Request body must be valid JSON'), }, + mapInput: ({ body }) => ({ workspaceId: body.workspaceId, path: body.path, source: 'api' }), + useCase: createKnowledgeFolder, + present: ({ folder }) => ({ data: { folder: toFolderPathView(folder, folder.path) } }), }) -export const PATCH = withPublicApiRouteHandler({ +export const PATCH = defineV2JsonRoute({ contract: v2RelocateKnowledgeFolderContract, - rateLimitEndpoint: 'knowledge', - handler: async ({ input, auth: { userId, rateLimit } }) => { - const { workspaceId, path, destinationPath } = input.body - const access = await resolveWorkspaceAccess(rateLimit, userId, workspaceId, 'write') - if (access) return v2WorkspaceAccessError(access) - const result = await relocateFolderByPath({ - resourceType: 'knowledge_base', - workspaceId, - userId, - path, - destinationPath, - }) - if (!result.success || !result.folder) { - return v2FolderPathMutationError(result.errorCode, result.error ?? 'Failed to move folder') - } - const index = await loadActiveFolderPathIndex(workspaceId, 'knowledge_base') - return v2Data({ folder: toV2PathFolder(result.folder, index, false) }, { rateLimit }) + auth: v2ApiKeyAuth, + operation: knowledgeOperations.relocateFolder, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + parseOptions: { + invalidJsonResponse: () => v2Error('BAD_REQUEST', 'Request body must be valid JSON'), }, + mapInput: ({ body }) => ({ + workspaceId: body.workspaceId, + path: body.path, + destinationPath: body.destinationPath, + source: 'api', + }), + useCase: relocateKnowledgeFolder, + present: ({ folder }) => ({ data: { folder: toFolderPathView(folder, folder.path) } }), }) -export const DELETE = withPublicApiRouteHandler({ +export const DELETE = defineV2JsonRoute({ contract: v2DeleteKnowledgeFolderContract, - rateLimitEndpoint: 'knowledge', - handler: async ({ input, auth: { userId, rateLimit } }) => { - const { workspaceId, path, recursive } = input.query - const access = await resolveWorkspaceAccess(rateLimit, userId, workspaceId, 'write') - if (access) return v2WorkspaceAccessError(access) - const result = await deleteFolderByPath({ - resourceType: 'knowledge_base', - workspaceId, - userId, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.deleteFolder, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + mapInput: ({ query }) => ({ + workspaceId: query.workspaceId, + path: query.path, + recursive: query.recursive, + source: 'api', + }), + useCase: deleteKnowledgeFolder, + present: ({ path, deletedItems }) => ({ + data: { path, - recursive, - }) - if (!result.success || !result.deletedItems) { - return v2FolderPathMutationError(result.errorCode, result.error ?? 'Failed to delete folder') - } - return v2Data( - { - path, - deleted: true as const, - deletedItems: { - folders: result.deletedItems.folders, - knowledgeBases: result.deletedItems.knowledgeBases ?? 0, - }, + deleted: true as const, + deletedItems: { + folders: deletedItems.folders, + knowledgeBases: deletedItems.knowledgeBases ?? 0, }, - { rateLimit } - ) - }, + }, + }), }) diff --git a/apps/sim/app/api/v2/knowledge/route.test.ts b/apps/sim/app/api/v2/knowledge/route.test.ts index bdd087b441e..9bdc6765bdc 100644 --- a/apps/sim/app/api/v2/knowledge/route.test.ts +++ b/apps/sim/app/api/v2/knowledge/route.test.ts @@ -1,157 +1,192 @@ /** * @vitest-environment node - * - * Public v2 knowledge-base list: the search/filter/sort convention reaching the - * lib rather than being applied over its result. */ import { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' const { + mockAuthenticate, + mockCheckPreAuth, mockCheckRateLimit, - mockResolveWorkspaceAccess, - mockGetKnowledgeBases, - mockLoadActiveFolderPathIndex, + mockList, + mockCreate, + mockPlatformCreated, + mockCapture, } = vi.hoisted(() => ({ + mockAuthenticate: vi.fn(), + mockCheckPreAuth: vi.fn(), mockCheckRateLimit: vi.fn(), - mockResolveWorkspaceAccess: vi.fn(), - mockGetKnowledgeBases: vi.fn(), - mockLoadActiveFolderPathIndex: vi.fn(), + mockList: vi.fn(), + mockCreate: vi.fn(), + mockPlatformCreated: vi.fn(), + mockCapture: vi.fn(), })) -vi.mock('@/app/api/v1/middleware', () => ({ - checkRateLimit: mockCheckRateLimit, - resolveWorkspaceAccess: mockResolveWorkspaceAccess, +vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => ({ + authenticateV2ApiKey: mockAuthenticate, + V2ApiKeyUnauthenticatedError: class V2ApiKeyUnauthenticatedError extends Error {}, })) -vi.mock('@/lib/knowledge/service', () => ({ - getKnowledgeBases: mockGetKnowledgeBases, +vi.mock('@/lib/core/rate-limiter', () => ({ + getRateLimit: () => ({ maxTokens: 100, refillRate: 100, refillIntervalMs: 60_000 }), + RateLimiter: class RateLimiter { + checkRateLimitDirect(...args: unknown[]) { + return mockCheckPreAuth(...args) + } + + checkRateLimitDirectOrThrow(...args: unknown[]) { + return mockCheckRateLimit(...args) + } + }, })) -vi.mock('@/lib/folders/queries', () => ({ - loadActiveFolderPathIndex: mockLoadActiveFolderPathIndex, -})) +vi.mock('@/app/api/v2/lib/gate', () => ({ v2ApiGateError: vi.fn().mockResolvedValue(null) })) -vi.mock('@/lib/knowledge/orchestration', () => ({ - performCreateKnowledgeBase: vi.fn(), +vi.mock('@/lib/knowledge/application/knowledge-bases', () => ({ + listKnowledgeBases: { operation: { id: 'knowledge.list' }, execute: mockList }, + createKnowledgeBase: { operation: { id: 'knowledge.create' }, execute: mockCreate }, })) -vi.mock('@/app/api/v2/lib/gate', () => ({ - v2ApiGateError: vi.fn().mockResolvedValue(null), +vi.mock('@/lib/core/telemetry', () => ({ + PlatformEvents: { knowledgeBaseCreated: mockPlatformCreated }, })) -import { GET } from '@/app/api/v2/knowledge/route' +vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: mockCapture })) -const WS = 'workspace-1' -const FOLDER_ID = 'fold_1' +import { GET, POST } from '@/app/api/v2/knowledge/route' +const WORKSPACE_ID = 'workspace-1' const RATE_LIMIT_OK = { allowed: true, - userId: 'user-1', - keyType: 'workspace', - limit: 100, remaining: 99, resetAt: new Date('2024-01-01T01:00:00Z'), + retryAfterMs: 0, } -/** What the route forwards for a bare `?workspaceId=` list. */ -const DEFAULT_LIST_ARGS = { - folderId: undefined, - search: undefined, - sortBy: 'createdAt', - sortOrder: 'asc', -} - -function buildKnowledgeBase(overrides: Record = {}) { +function buildKnowledgeBase() { return { - id: 'kb_1', + id: 'kb-1', userId: 'user-1', name: 'Support docs', description: null, tokenCount: 0, embeddingModel: 'text-embedding-3-small', embeddingDimension: 1536, - chunkingConfig: { maxSize: 1024, minSize: 1, overlap: 200 }, - workspaceId: WS, + chunkingConfig: { maxSize: 1024, minSize: 100, overlap: 200 }, + workspaceId: WORKSPACE_ID, folderId: null, docCount: 2, + connectorTypes: ['notion'], createdAt: new Date('2024-01-01T00:00:00Z'), updatedAt: new Date('2024-01-02T00:00:00Z'), deletedAt: null, - ...overrides, } } -const callList = (query: string) => - GET(new NextRequest(`http://localhost:3000/api/v2/knowledge?${query}`)) - -describe('GET /api/v2/knowledge', () => { +describe('/api/v2/knowledge route composition', () => { beforeEach(() => { vi.clearAllMocks() + mockCheckPreAuth.mockResolvedValue(RATE_LIMIT_OK) mockCheckRateLimit.mockResolvedValue(RATE_LIMIT_OK) - mockResolveWorkspaceAccess.mockResolvedValue(null) - mockGetKnowledgeBases.mockResolvedValue([buildKnowledgeBase()]) - mockLoadActiveFolderPathIndex.mockResolvedValue({ - rowById: new Map([['fold_1', { id: 'fold_1', name: 'Support', parentId: null }]]), - pathById: new Map([['fold_1', '/Support']]), - idByPath: new Map([['/Support', 'fold_1']]), + mockAuthenticate.mockResolvedValue({ + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + rolloutUserId: 'user-1', + rateLimitSubjectIds: ['api-key:key-1', 'user:user-1'], + rateLimitSubscription: null, + keyType: 'personal', + }) + mockList.mockResolvedValue({ + knowledgeBases: [{ knowledgeBase: buildKnowledgeBase(), folderPath: '/' }], }) + mockCreate.mockResolvedValue({ knowledgeBase: buildKnowledgeBase(), folderPath: '/' }) }) - it('forwards search, folder, and sort into the query rather than filtering the result', async () => { - const res = await callList( - `workspaceId=${WS}&search=support&folderPath=${encodeURIComponent('/Support')}&sortBy=name&sortOrder=desc` + it('delegates the bounded list query with the authenticated principal', async () => { + const request = new NextRequest( + `http://localhost/api/v2/knowledge?workspaceId=${WORKSPACE_ID}&search=support&folderPath=%2F&sortBy=name&sortOrder=desc`, + { headers: { 'x-api-key': 'secret' } } ) - expect(res.status).toBe(200) - expect(mockGetKnowledgeBases).toHaveBeenCalledWith('user-1', WS, 'active', { - folderId: FOLDER_ID, - search: 'support', - sortBy: 'name', - sortOrder: 'desc', + const response = await GET(request) + + expect(response.status).toBe(200) + expect(mockList).toHaveBeenCalledWith({ + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + input: { + workspaceId: WORKSPACE_ID, + folderPath: '/', + search: 'support', + sortBy: 'name', + sortOrder: 'desc', + }, + request, }) - }) - - it('defaults to the createdAt ordering when no sort is requested', async () => { - await callList(`workspaceId=${WS}`) - - expect(mockGetKnowledgeBases).toHaveBeenCalledWith('user-1', WS, 'active', DEFAULT_LIST_ARGS) - }) - - it('treats folderPath=/ as root-only while omission lists every folder', async () => { - await callList(`workspaceId=${WS}&folderPath=%2F`) - - expect(mockGetKnowledgeBases).toHaveBeenCalledWith('user-1', WS, 'active', { - ...DEFAULT_LIST_ARGS, - folderId: null, + expect(await response.json()).toEqual({ + data: [ + expect.objectContaining({ + id: 'kb-1', + folderPath: '/', + connectorTypes: ['notion'], + createdAt: '2024-01-01T00:00:00.000Z', + }), + ], + nextCursor: null, }) }) - it('400s on a sort field outside the enum instead of letting it reach the query', async () => { - const res = await callList(`workspaceId=${WS}&sortBy=name);--`) - - expect(res.status).toBe(400) - expect((await res.json()).error.code).toBe('BAD_REQUEST') - expect(mockGetKnowledgeBases).not.toHaveBeenCalled() - }) - - it('400s on a sort direction outside the enum', async () => { - const res = await callList(`workspaceId=${WS}&sortOrder=sideways`) + it('returns 201 and keeps human analytics on the personal-key actor', async () => { + const request = new NextRequest('http://localhost/api/v2/knowledge', { + method: 'POST', + headers: { 'content-type': 'application/json', 'x-api-key': 'secret' }, + body: JSON.stringify({ workspaceId: WORKSPACE_ID, name: 'Support docs' }), + }) - expect(res.status).toBe(400) - expect(mockGetKnowledgeBases).not.toHaveBeenCalled() + const response = await POST(request) + + expect(response.status).toBe(201) + expect(mockCreate).toHaveBeenCalledWith({ + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + input: { + workspaceId: WORKSPACE_ID, + name: 'Support docs', + description: undefined, + chunkingConfig: { maxSize: 1024, minSize: 100, overlap: 200 }, + folderPath: undefined, + source: 'api', + }, + request, + }) + expect(mockPlatformCreated).toHaveBeenCalledWith({ + knowledgeBaseId: 'kb-1', + name: 'Support docs', + workspaceId: WORKSPACE_ID, + }) + expect(mockCapture).toHaveBeenCalledWith( + 'user-1', + 'knowledge_base_created', + expect.objectContaining({ workspace_id: WORKSPACE_ID }), + expect.any(Object) + ) }) - it('400s on an empty search rather than treating it as unsearched', async () => { - const res = await callList(`workspaceId=${WS}&search=`) - - expect(res.status).toBe(400) - expect(mockGetKnowledgeBases).not.toHaveBeenCalled() - }) + it('does not attribute workspace-key creation analytics to a billing owner', async () => { + mockAuthenticate.mockResolvedValue({ + principal: { kind: 'workspace_api_key', workspaceId: WORKSPACE_ID, keyId: 'key-2' }, + rolloutUserId: 'billing-owner', + rateLimitSubjectIds: ['api-key:key-2', `workspace:${WORKSPACE_ID}`], + rateLimitSubscription: null, + keyType: 'workspace', + }) + const request = new NextRequest('http://localhost/api/v2/knowledge', { + method: 'POST', + headers: { 'content-type': 'application/json', 'x-api-key': 'secret' }, + body: JSON.stringify({ workspaceId: WORKSPACE_ID, name: 'Support docs' }), + }) - it('terminates pagination with a filter applied', async () => { - const res = await callList(`workspaceId=${WS}&search=support`) + const response = await POST(request) - expect((await res.json()).nextCursor).toBeNull() + expect(response.status).toBe(201) + expect(mockPlatformCreated).toHaveBeenCalledOnce() + expect(mockCapture).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/app/api/v2/knowledge/route.ts b/apps/sim/app/api/v2/knowledge/route.ts index 8b20a6b7d3d..900cbb06ee3 100644 --- a/apps/sim/app/api/v2/knowledge/route.ts +++ b/apps/sim/app/api/v2/knowledge/route.ts @@ -2,101 +2,122 @@ import { v2CreateKnowledgeBaseContract, v2ListKnowledgeBasesContract, } from '@/lib/api/contracts/v2/knowledge' -import { loadActiveFolderPathIndex } from '@/lib/folders/queries' -import { performCreateKnowledgeBase } from '@/lib/knowledge/orchestration' -import { getKnowledgeBases } from '@/lib/knowledge/service' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' -import { formatKnowledgeBase } from '@/app/api/v1/knowledge/utils' -import { resolveWorkspaceAccess } from '@/app/api/v1/middleware' import { - folderPathForId, - resolveFolderPathId, - resolveFolderPathIdentity, -} from '@/app/api/v2/lib/folders' + defineV2JsonRoute, + v2ApiKeyAuth, + v2OrchestrationErrorPolicy, + v2RateLimits, +} from '@/lib/api/server/routes' +import { PlatformEvents } from '@/lib/core/telemetry' import { - v2CursorList, - v2Data, - v2Error, - v2ErrorForOrchestration, - v2WorkspaceAccessError, -} from '@/app/api/v2/lib/response' + createKnowledgeBase, + listKnowledgeBases, +} from '@/lib/knowledge/application/knowledge-bases' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types' +import { captureServerEvent } from '@/lib/posthog/server' +import { v2Error } from '@/app/api/v2/lib/response' export const dynamic = 'force-dynamic' export const revalidate = 0 +function toV2KnowledgeBase(knowledgeBase: KnowledgeBaseWithCounts, folderPath: string) { + return { + id: knowledgeBase.id, + name: knowledgeBase.name, + description: knowledgeBase.description, + tokenCount: knowledgeBase.tokenCount, + embeddingModel: knowledgeBase.embeddingModel, + embeddingDimension: knowledgeBase.embeddingDimension, + chunkingConfig: { + maxSize: knowledgeBase.chunkingConfig.maxSize, + minSize: knowledgeBase.chunkingConfig.minSize, + overlap: knowledgeBase.chunkingConfig.overlap, + strategy: knowledgeBase.chunkingConfig.strategy, + strategyOptions: knowledgeBase.chunkingConfig.strategyOptions + ? { + pattern: knowledgeBase.chunkingConfig.strategyOptions.pattern, + separators: knowledgeBase.chunkingConfig.strategyOptions.separators, + recipe: knowledgeBase.chunkingConfig.strategyOptions.recipe, + strictBoundaries: knowledgeBase.chunkingConfig.strategyOptions.strictBoundaries, + } + : undefined, + }, + docCount: knowledgeBase.docCount, + connectorTypes: knowledgeBase.connectorTypes, + createdAt: knowledgeBase.createdAt.toISOString(), + updatedAt: knowledgeBase.updatedAt.toISOString(), + folderPath, + } +} + /** GET /api/v2/knowledge — List knowledge bases in a workspace. */ -export const GET = withPublicApiRouteHandler({ +export const GET = defineV2JsonRoute({ contract: v2ListKnowledgeBasesContract, - rateLimitEndpoint: 'knowledge', - handler: async ({ input, auth: { userId, rateLimit } }) => { - const { workspaceId, folderPath, search, sortBy, sortOrder } = input.query - - const access = await resolveWorkspaceAccess(rateLimit, userId, workspaceId, 'read') - if (access) return v2WorkspaceAccessError(access) - - const folderIndex = await loadActiveFolderPathIndex(workspaceId, 'knowledge_base') - const folderId = - folderPath === undefined ? undefined : resolveFolderPathId(folderIndex, folderPath) - if (folderPath !== undefined && folderId === undefined) { - return v2Error('NOT_FOUND', 'Folder not found') - } - - const knowledgeBases = await getKnowledgeBases(userId, workspaceId, 'active', { - folderId, - search, - sortBy, - sortOrder, - }) - const items = knowledgeBases.map((knowledgeBase) => ({ - ...formatKnowledgeBase(knowledgeBase), - folderPath: folderPathForId(folderIndex, knowledgeBase.folderId), - })) - - // `getKnowledgeBases` returns the full bounded workspace set → single page. - return v2CursorList(items, null, { rateLimit }) - }, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.list, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + mapInput: ({ query }) => ({ + workspaceId: query.workspaceId, + folderPath: query.folderPath, + search: query.search, + sortBy: query.sortBy, + sortOrder: query.sortOrder, + }), + useCase: listKnowledgeBases, + present: ({ knowledgeBases }) => ({ + data: knowledgeBases.map(({ knowledgeBase, folderPath }) => + toV2KnowledgeBase(knowledgeBase, folderPath) + ), + nextCursor: null, + }), }) /** POST /api/v2/knowledge — Create a new knowledge base. */ -export const POST = withPublicApiRouteHandler({ +export const POST = defineV2JsonRoute({ contract: v2CreateKnowledgeBaseContract, - rateLimitEndpoint: 'knowledge', - handler: async ({ request, input, auth: { requestId, userId, rateLimit } }) => { - const { workspaceId, name, description, chunkingConfig, folderPath } = input.body - - const access = await resolveWorkspaceAccess(rateLimit, userId, workspaceId, 'write') - if (access) return v2WorkspaceAccessError(access) - - const resolution = await resolveFolderPathIdentity({ - workspaceId, - resourceType: 'knowledge_base', - path: folderPath ?? '/', - }) - if (!resolution.found) return v2Error('NOT_FOUND', 'Folder not found') - - const outcome = await performCreateKnowledgeBase({ - userId, - source: 'api', - workspaceId, - name, - description, - chunkingConfig, - folderId: resolution.folderId, - requestId, - request, + auth: v2ApiKeyAuth, + operation: knowledgeOperations.create, + rateLimit: v2RateLimits.publicApi, + errorPolicy: v2OrchestrationErrorPolicy, + parseOptions: { + invalidJsonResponse: () => v2Error('BAD_REQUEST', 'Request body must be valid JSON'), + }, + mapInput: ({ body }) => ({ + workspaceId: body.workspaceId, + name: body.name, + description: body.description, + chunkingConfig: body.chunkingConfig, + folderPath: body.folderPath, + source: 'api', + }), + useCase: createKnowledgeBase, + onSuccess: ({ principal, result: { knowledgeBase } }) => { + PlatformEvents.knowledgeBaseCreated({ + knowledgeBaseId: knowledgeBase.id, + name: knowledgeBase.name, + workspaceId: knowledgeBase.workspaceId ?? undefined, }) - if (!outcome.success) { - return v2ErrorForOrchestration(outcome.errorCode, outcome.error) - } - - return v2Data( - { - knowledgeBase: { - ...formatKnowledgeBase(outcome.knowledgeBase), - folderPath: folderPathForId(resolution.index, outcome.knowledgeBase.folderId), + if (principal.kind === 'personal_api_key') { + captureServerEvent( + principal.userId, + 'knowledge_base_created', + { + knowledge_base_id: knowledgeBase.id, + workspace_id: knowledgeBase.workspaceId ?? '', + name: knowledgeBase.name, }, - }, - { rateLimit, status: 201 } - ) + { + ...(knowledgeBase.workspaceId + ? { groups: { workspace: knowledgeBase.workspaceId } } + : {}), + setOnce: { first_kb_created_at: new Date().toISOString() }, + } + ) + } }, + present: ({ knowledgeBase, folderPath }) => ({ + data: { knowledgeBase: toV2KnowledgeBase(knowledgeBase, folderPath) }, + }), }) diff --git a/apps/sim/app/api/v2/knowledge/search/route.test.ts b/apps/sim/app/api/v2/knowledge/search/route.test.ts new file mode 100644 index 00000000000..462518ba38f --- /dev/null +++ b/apps/sim/app/api/v2/knowledge/search/route.test.ts @@ -0,0 +1,145 @@ +/** + * @vitest-environment node + */ +import { NextRequest } from 'next/server' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockAuthenticate, mockCheckPreAuth, mockCheckRateLimit, mockSearch } = vi.hoisted(() => ({ + mockAuthenticate: vi.fn(), + mockCheckPreAuth: vi.fn(), + mockCheckRateLimit: vi.fn(), + mockSearch: vi.fn(), +})) + +vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => ({ + authenticateV2ApiKey: mockAuthenticate, + V2ApiKeyUnauthenticatedError: class V2ApiKeyUnauthenticatedError extends Error {}, +})) + +vi.mock('@/lib/core/rate-limiter', () => ({ + getRateLimit: () => ({ maxTokens: 100, refillRate: 100, refillIntervalMs: 60_000 }), + RateLimiter: class RateLimiter { + checkRateLimitDirect(...args: unknown[]) { + return mockCheckPreAuth(...args) + } + + checkRateLimitDirectOrThrow(...args: unknown[]) { + return mockCheckRateLimit(...args) + } + }, +})) + +vi.mock('@/app/api/v2/lib/gate', () => ({ v2ApiGateError: vi.fn().mockResolvedValue(null) })) + +vi.mock('@/lib/knowledge/application/search', () => ({ + searchKnowledge: { operation: { id: 'knowledge.search' }, execute: mockSearch }, +})) + +import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing' +import { POST } from '@/app/api/v2/knowledge/search/route' + +const WORKSPACE_ID = 'workspace-1' +const PRINCIPAL = { kind: 'workspace_api_key', workspaceId: WORKSPACE_ID, keyId: 'key-1' } as const +const RATE_LIMIT_OK = { + allowed: true, + remaining: 99, + resetAt: new Date('2024-01-01T01:00:00Z'), + retryAfterMs: 0, +} + +function buildRequest(body: string) { + return new NextRequest('http://localhost/api/v2/knowledge/search', { + method: 'POST', + headers: { 'content-type': 'application/json', 'x-api-key': 'secret' }, + body, + }) +} + +describe('POST /api/v2/knowledge/search', () => { + beforeEach(() => { + vi.clearAllMocks() + mockCheckPreAuth.mockResolvedValue(RATE_LIMIT_OK) + mockCheckRateLimit.mockResolvedValue(RATE_LIMIT_OK) + mockAuthenticate.mockResolvedValue({ + principal: PRINCIPAL, + rolloutUserId: 'billing-owner', + rateLimitSubjectIds: ['api-key:key-1', `workspace:${WORKSPACE_ID}`], + rateLimitSubscription: null, + keyType: 'workspace', + }) + mockSearch.mockResolvedValue({ + results: [ + { + documentId: 'doc-1', + documentName: 'support.txt', + sourceUrl: null, + content: 'hello', + chunkIndex: 0, + metadata: {}, + similarity: 0.9, + }, + ], + query: 'hello', + knowledgeBaseIds: ['kb-1'], + topK: 10, + totalResults: 1, + }) + }) + + it('delegates normalized IDs through the semantic operation', async () => { + const request = buildRequest( + JSON.stringify({ + workspaceId: WORKSPACE_ID, + knowledgeBaseIds: 'kb-1', + query: 'hello', + topK: 10, + }) + ) + + const response = await POST(request) + + expect(response.status).toBe(200) + expect(mockSearch).toHaveBeenCalledWith({ + principal: PRINCIPAL, + input: { + workspaceId: WORKSPACE_ID, + knowledgeBaseIds: ['kb-1'], + query: 'hello', + topK: 10, + tagFilters: undefined, + }, + request, + }) + expect(await response.json()).toEqual({ + data: expect.objectContaining({ knowledgeBaseIds: ['kb-1'], totalResults: 1 }), + }) + }) + + it('authenticates before rejecting malformed JSON', async () => { + const response = await POST(buildRequest('{')) + + expect(response.status).toBe(400) + expect(mockAuthenticate).toHaveBeenCalledOnce() + expect(mockSearch).not.toHaveBeenCalled() + }) + + it('maps usage failures without exposing infrastructure details', async () => { + mockSearch.mockRejectedValue(new KnowledgeUsageLimitExceededError('Upgrade required')) + + const response = await POST( + buildRequest( + JSON.stringify({ + workspaceId: WORKSPACE_ID, + knowledgeBaseIds: ['kb-1'], + query: 'hello', + topK: 10, + }) + ) + ) + + expect(response.status).toBe(402) + expect(await response.json()).toEqual({ + error: { code: 'USAGE_LIMIT_EXCEEDED', message: 'Upgrade required' }, + }) + }) +}) diff --git a/apps/sim/app/api/v2/knowledge/search/route.ts b/apps/sim/app/api/v2/knowledge/search/route.ts index dc6f3fdc829..3aa8d98f158 100644 --- a/apps/sim/app/api/v2/knowledge/search/route.ts +++ b/apps/sim/app/api/v2/knowledge/search/route.ts @@ -1,271 +1,74 @@ -import { - type V2KnowledgeSearchResult, - v2SearchKnowledgeContract, -} from '@/lib/api/contracts/v2/knowledge' -import { isZodError } from '@/lib/api/server' -import { - checkAttributedUsageLimits, - resolveBillingAttribution, - resolveSystemBillingAttribution, -} from '@/lib/billing/core/billing-attribution' -import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants' -import { recordSearchEmbeddingUsage } from '@/lib/knowledge/embeddings' -import { - executeKnowledgeSearch, - generateSearchEmbedding, - getDocumentMetadataByIds, - type SearchResult, -} from '@/lib/knowledge/search/queries' -import { getDocumentTagDefinitions } from '@/lib/knowledge/tags/service' -import { buildUndefinedTagsError, validateTagValue } from '@/lib/knowledge/tags/utils' -import type { StructuredFilter } from '@/lib/knowledge/types' -import { checkKnowledgeBaseAccess, type KnowledgeBaseAccessResult } from '@/app/api/knowledge/utils' -import { withPublicApiRouteHandler } from '@/app/api/public-api-route-handler' -import { resolveWorkspaceAccess } from '@/app/api/v1/middleware' -import { - v2Data, - v2Error, - v2ValidationError, - v2WorkspaceAccessError, -} from '@/app/api/v2/lib/response' +import type { NextRequest } from 'next/server' +import { NextResponse } from 'next/server' +import { v2SearchKnowledgeContract } from '@/lib/api/contracts/v2/knowledge' +import { parseRequest } from '@/lib/api/server' +import { v2ApiKeyAuth, v2OrchestrationErrorPolicy, v2RateLimits } from '@/lib/api/server/routes' +import type { JsonRouteContext } from '@/lib/api/server/routes/types' +import { admitV2Request, V2RouteInfrastructureError } from '@/lib/api/server/routes/v2-json-route' +import { withRouteHandler } from '@/lib/core/utils/with-route-handler' +import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { searchKnowledge } from '@/lib/knowledge/application/search' +import { v2Error, v2ValidationError } from '@/app/api/v2/lib/response' export const dynamic = 'force-dynamic' export const revalidate = 0 /** POST /api/v2/knowledge/search — Vector / tag search across knowledge bases. */ -export const POST = withPublicApiRouteHandler({ - contract: v2SearchKnowledgeContract, - rateLimitEndpoint: 'knowledge-search', - handler: async ({ input, auth: { requestId, userId, rateLimit } }) => { - try { - const { workspaceId, topK, query, tagFilters } = input.body - - const access = await resolveWorkspaceAccess(rateLimit, userId, workspaceId, 'read') - if (access) return v2WorkspaceAccessError(access) - - /** - * A query incurs hosted embedding (+ optional rerank) cost — gate the - * actor's usage before spending; tag-only search is free. Workspace keys - * resolve their system actor and immutable payer from one workspace read. - */ - const hasBillableQuery = Boolean(query?.trim()) - const billingAttribution = hasBillableQuery - ? rateLimit.keyType === 'workspace' - ? await resolveSystemBillingAttribution(workspaceId) - : await resolveBillingAttribution({ actorUserId: userId, workspaceId }) - : undefined - const billingActorUserId = billingAttribution?.actorUserId ?? userId - if (billingAttribution) { - const usage = await checkAttributedUsageLimits(billingAttribution) - if (usage.isExceeded) { - return v2Error( - 'USAGE_LIMIT_EXCEEDED', - usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' - ) - } - } - - const knowledgeBaseIds = Array.isArray(input.body.knowledgeBaseIds) - ? input.body.knowledgeBaseIds - : [input.body.knowledgeBaseIds] - - const accessChecks = await Promise.all( - knowledgeBaseIds.map((kbId) => checkKnowledgeBaseAccess(kbId, userId)) +export const POST = withRouteHandler( + async (request: NextRequest, context) => { + if (request.method !== v2SearchKnowledgeContract.method) { + throw new Error( + `Route received ${request.method} for ${v2SearchKnowledgeContract.method} contract ${v2SearchKnowledgeContract.path}` ) - const accessibleKbs = accessChecks - .filter( - (ac): ac is KnowledgeBaseAccessResult => - ac.hasAccess === true && ac.knowledgeBase.workspaceId === workspaceId - ) - .map((ac) => ac.knowledgeBase) - const accessibleKbIds = accessibleKbs.map((kb) => kb.id) - - if (accessibleKbIds.length === 0) { - return v2Error('NOT_FOUND', 'Knowledge base not found or access denied') - } - - const inaccessibleKbIds = knowledgeBaseIds.filter((id) => !accessibleKbIds.includes(id)) - if (inaccessibleKbIds.length > 0) { - return v2Error( - 'NOT_FOUND', - `Knowledge bases not found or access denied: ${inaccessibleKbIds.join(', ')}` - ) - } - - let structuredFilters: StructuredFilter[] = [] - const tagDefsCache = new Map>>() - - if (tagFilters && tagFilters.length > 0 && accessibleKbIds.length > 1) { - return v2Error( - 'BAD_REQUEST', - 'Tag filters are only supported when searching a single knowledge base' - ) - } - - if (tagFilters && tagFilters.length > 0 && accessibleKbIds.length > 0) { - const kbId = accessibleKbIds[0] - const tagDefs = await getDocumentTagDefinitions(kbId) - tagDefsCache.set(kbId, tagDefs) - - const displayNameToTagDef: Record = {} - tagDefs.forEach((def) => { - displayNameToTagDef[def.displayName] = { - tagSlot: def.tagSlot, - fieldType: def.fieldType, - } - }) - - const undefinedTags: string[] = [] - const typeErrors: string[] = [] - - for (const filter of tagFilters) { - const tagDef = displayNameToTagDef[filter.tagName] - if (!tagDef) { - undefinedTags.push(filter.tagName) - continue - } - const validationError = validateTagValue( - filter.tagName, - String(filter.value), - tagDef.fieldType - ) - if (validationError) { - typeErrors.push(validationError) - } - } - - if (undefinedTags.length > 0 || typeErrors.length > 0) { - const errorParts: string[] = [] - if (undefinedTags.length > 0) { - errorParts.push(buildUndefinedTagsError(undefinedTags)) - } - if (typeErrors.length > 0) { - errorParts.push(...typeErrors) - } - return v2Error('BAD_REQUEST', errorParts.join('\n')) - } - - structuredFilters = tagFilters.map((filter) => { - const tagDef = displayNameToTagDef[filter.tagName]! - return { - tagSlot: tagDef.tagSlot, - fieldType: tagDef.fieldType, - operator: filter.operator, - value: filter.value, - valueTo: filter.valueTo, - } - }) - } - - const hasQuery = Boolean(query && query.trim().length > 0) - const hasFilters = structuredFilters.length > 0 - - const embeddingModels = Array.from(new Set(accessibleKbs.map((kb) => kb.embeddingModel))) - if (hasQuery && embeddingModels.length > 1) { - return v2Error( - 'BAD_REQUEST', - 'Selected knowledge bases use different embedding models and cannot be searched together. Search them separately.' - ) - } - const queryEmbeddingModel = embeddingModels[0] - - if (!hasQuery && !hasFilters) { - return v2Error('BAD_REQUEST', 'Either query or tagFilters must be provided') - } - - let queryEmbeddingIsBYOK: boolean | null = null - let queryVector: string | undefined - - if (hasQuery) { - const queryEmbeddingResult = await generateSearchEmbedding( - query!, - queryEmbeddingModel, - workspaceId - ) - queryEmbeddingIsBYOK = queryEmbeddingResult.isBYOK - queryVector = JSON.stringify(queryEmbeddingResult.embedding) - } - - const results: SearchResult[] = await executeKnowledgeSearch({ - knowledgeBaseIds: accessibleKbIds, - topK, - searchMode: 'vector', - query, - queryVector, - structuredFilters, - }) - - if (queryEmbeddingIsBYOK !== null) { - await recordSearchEmbeddingUsage({ - userId: billingActorUserId, - workspaceId, - embeddingModel: queryEmbeddingModel, - query: query!, - isBYOK: queryEmbeddingIsBYOK, - sourceReference: `v2-kb-search:${requestId}`, - billingAttribution, - }) - } + } - const tagDefsResults = await Promise.all( - accessibleKbIds.map(async (kbId) => { - try { - const tagDefs = tagDefsCache.get(kbId) ?? (await getDocumentTagDefinitions(kbId)) - const map: Record = {} - tagDefs.forEach((def) => { - map[def.tagSlot] = def.displayName - }) - return { kbId, map } - } catch { - return { kbId, map: {} as Record } - } - }) - ) - const tagDefinitionsMap: Record> = {} - tagDefsResults.forEach(({ kbId, map }) => { - tagDefinitionsMap[kbId] = map + const admission = await admitV2Request( + request, + knowledgeOperations.search, + v2ApiKeyAuth, + v2RateLimits.publicApi + ) + if (!admission.success) return admission.response + + const parsed = await parseRequest(v2SearchKnowledgeContract, request, context ?? {}, { + validationErrorResponse: v2ValidationError, + invalidJsonResponse: () => v2Error('BAD_REQUEST', 'Request body must be valid JSON'), + }) + if (!parsed.success) return parsed.response + + const { body } = parsed.data + try { + const result = await searchKnowledge.execute({ + principal: admission.auth.principal, + input: { + workspaceId: body.workspaceId, + knowledgeBaseIds: Array.isArray(body.knowledgeBaseIds) + ? body.knowledgeBaseIds + : [body.knowledgeBaseIds], + query: body.query, + topK: body.topK, + tagFilters: body.tagFilters, + }, + request, }) - - const documentIds = results.map((r) => r.documentId) - const documentMetadataMap = await getDocumentMetadataByIds(documentIds) - - const searchResults: V2KnowledgeSearchResult[] = results.map((result) => { - const kbTagMap = tagDefinitionsMap[result.knowledgeBaseId] || {} - const metadata: Record = {} - - ALL_TAG_SLOTS.forEach((slot) => { - const tagValue = result[slot as keyof SearchResult] - if (tagValue !== null && tagValue !== undefined) { - const displayName = kbTagMap[slot] || slot - metadata[displayName] = tagValue - } - }) - - const docMeta = documentMetadataMap[result.documentId] - return { - documentId: result.documentId, - documentName: docMeta?.filename ?? null, - sourceUrl: docMeta?.sourceUrl ?? null, - content: result.content, - chunkIndex: result.chunkIndex, - metadata, - similarity: hasQuery ? 1 - result.distance : 1, - } + const responseBody = v2SearchKnowledgeContract.response.schema.parse({ data: result }) + return NextResponse.json(responseBody, { + headers: { 'Cache-Control': 'private, no-store' }, }) - - return v2Data( - { - results: searchResults, - query: query || '', - knowledgeBaseIds: accessibleKbIds, - topK, - totalResults: results.length, - }, - { rateLimit } - ) } catch (error) { - if (isZodError(error)) return v2ValidationError(error) + if (error instanceof KnowledgeUsageLimitExceededError) { + return v2Error('USAGE_LIMIT_EXCEEDED', error.message) + } + const response = v2OrchestrationErrorPolicy.render(error) + if (response) return response throw error } }, -}) + { + unhandledErrorResponse: ({ error }) => + error instanceof V2RouteInfrastructureError + ? v2Error('SERVICE_UNAVAILABLE', 'Service temporarily unavailable') + : v2Error('INTERNAL_ERROR', 'Internal server error'), + } +) diff --git a/apps/sim/lib/api/contracts/knowledge/upload-sessions.ts b/apps/sim/lib/api/contracts/knowledge/upload-sessions.ts index 01340182786..b37ad76428b 100644 --- a/apps/sim/lib/api/contracts/knowledge/upload-sessions.ts +++ b/apps/sim/lib/api/contracts/knowledge/upload-sessions.ts @@ -21,6 +21,7 @@ export const createKnowledgeDocumentUploadContract = defineRouteContract({ response: { mode: 'json', schema: v2DataResponse(v2CreateKnowledgeDocumentUploadDataSchema), + status: 201, }, }) diff --git a/apps/sim/lib/api/contracts/v2/__tests__/knowledge.test.ts b/apps/sim/lib/api/contracts/v2/__tests__/knowledge.test.ts new file mode 100644 index 00000000000..23960193484 --- /dev/null +++ b/apps/sim/lib/api/contracts/v2/__tests__/knowledge.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from 'vitest' +import { + v2CreateKnowledgeBaseContract, + v2CreateKnowledgeDocumentUploadContract, + v2CreateKnowledgeFolderContract, + v2SearchKnowledgeContract, + v2UploadKnowledgeDocumentContract, +} from '@/lib/api/contracts/v2/knowledge' + +describe('v2 knowledge contracts', () => { + it('declares 201 for every resource-creation response', () => { + expect(v2CreateKnowledgeBaseContract.response.status).toBe(201) + expect(v2CreateKnowledgeFolderContract.response.status).toBe(201) + expect(v2UploadKnowledgeDocumentContract.response.status).toBe(201) + expect(v2CreateKnowledgeDocumentUploadContract.response.status).toBe(201) + }) + + it('preserves knowledge search bounds', () => { + const valid = v2SearchKnowledgeContract.body?.safeParse({ + workspaceId: 'workspace-1', + knowledgeBaseIds: Array.from({ length: 20 }, (_, index) => `kb-${index}`), + query: 'support', + topK: 100, + }) + const tooManyKnowledgeBases = v2SearchKnowledgeContract.body?.safeParse({ + workspaceId: 'workspace-1', + knowledgeBaseIds: Array.from({ length: 21 }, (_, index) => `kb-${index}`), + query: 'support', + topK: 100, + }) + const excessiveTopK = v2SearchKnowledgeContract.body?.safeParse({ + workspaceId: 'workspace-1', + knowledgeBaseIds: ['kb-1'], + query: 'support', + topK: 101, + }) + + expect(valid?.success).toBe(true) + expect(tooManyKnowledgeBases?.success).toBe(false) + expect(excessiveTopK?.success).toBe(false) + }) +}) diff --git a/apps/sim/lib/api/contracts/v2/knowledge.ts b/apps/sim/lib/api/contracts/v2/knowledge.ts index 5ffc9ee7580..4ee31c3b933 100644 --- a/apps/sim/lib/api/contracts/v2/knowledge.ts +++ b/apps/sim/lib/api/contracts/v2/knowledge.ts @@ -307,6 +307,7 @@ export const v2CreateKnowledgeBaseContract = defineRouteContract({ response: { mode: 'json', schema: v2DataResponse(v2KnowledgeBaseDataSchema), + status: 201, }, }) @@ -362,7 +363,7 @@ export const v2CreateKnowledgeFolderContract = defineRouteContract({ method: 'POST', path: '/api/v2/knowledge/folders', body: v2CreateFolderBodySchema, - response: { mode: 'json', schema: v2DataResponse(v2KnowledgeFolderDataSchema) }, + response: { mode: 'json', schema: v2DataResponse(v2KnowledgeFolderDataSchema), status: 201 }, }) export const v2RelocateKnowledgeFolderContract = defineRouteContract({ @@ -418,6 +419,7 @@ export const v2UploadKnowledgeDocumentContract = defineRouteContract({ response: { mode: 'json', schema: v2DataResponse(v2KnowledgeDocumentSummaryDataSchema), + status: 201, }, }) @@ -429,6 +431,7 @@ export const v2CreateKnowledgeDocumentUploadContract = defineRouteContract({ response: { mode: 'json', schema: v2DataResponse(v2CreateKnowledgeDocumentUploadDataSchema), + status: 201, }, }) diff --git a/apps/sim/lib/copilot/application/execute-knowledge-use-case.ts b/apps/sim/lib/copilot/application/execute-knowledge-use-case.ts new file mode 100644 index 00000000000..f9f9fb162f2 --- /dev/null +++ b/apps/sim/lib/copilot/application/execute-knowledge-use-case.ts @@ -0,0 +1,65 @@ +import type { DelegatedPrincipal } from '@sim/auth/principal' +import type { OperationUseCase } from '@/lib/core/application' +import { asOrchestrationError } from '@/lib/core/orchestration/types' +import { createKnowledgeDelegatedPrincipal } from '@/lib/knowledge/application/delegated-principal' +import { + type KnowledgeOperation, + knowledgeOperations, +} from '@/lib/knowledge/application/operations' + +export interface CopilotKnowledgeDelegationContext { + userId: string + workspaceId?: string + chatId?: string + executionId?: string + toolCallId?: string + copilotToolExecution?: boolean +} + +const registeredKnowledgeOperationIds = new Set( + Object.values(knowledgeOperations).map((operation) => operation.id) +) + +/** Normalizes immutable Copilot execution identity into a knowledge delegation. */ +export function resolveCopilotKnowledgePrincipal( + context: CopilotKnowledgeDelegationContext | undefined +): DelegatedPrincipal { + if (!context) throw new Error('Knowledge delegation requires a Copilot execution context') + if (!context.copilotToolExecution) { + throw new Error('Knowledge delegation requires a trusted Copilot execution context') + } + if (!context.userId) throw new Error('Knowledge delegation requires an authenticated user ID') + if (!context.workspaceId) throw new Error('Knowledge delegation requires a workspace ID') + if (!context.toolCallId) throw new Error('Knowledge delegation requires a tool call ID') + + return createKnowledgeDelegatedPrincipal({ + serviceId: 'copilot', + subjectUserId: context.userId, + workspaceId: context.workspaceId, + delegationId: context.toolCallId, + chatId: context.chatId, + executionId: context.executionId, + }) +} + +/** Enters a registered knowledge application use case with trusted Copilot identity. */ +export function executeCopilotKnowledgeUseCase( + context: CopilotKnowledgeDelegationContext | undefined, + useCase: OperationUseCase, + input: I +): Promise { + if (!registeredKnowledgeOperationIds.has(useCase.operation.id)) { + throw new Error(`Unregistered Copilot knowledge operation: ${useCase.operation.id}`) + } + return useCase.execute({ principal: resolveCopilotKnowledgePrincipal(context), input }) +} + +/** Projects only caller-actionable application errors into a Copilot result. */ +export function messageForCopilotKnowledgeError( + error: unknown, + fallback = 'Knowledge operation failed' +): string { + const classified = asOrchestrationError(error) + if (classified && classified.code !== 'internal') return classified.message + return fallback +} diff --git a/apps/sim/lib/copilot/tools/handlers/resources.test.ts b/apps/sim/lib/copilot/tools/handlers/resources.test.ts index f470d47e6da..6b5f8595ea2 100644 --- a/apps/sim/lib/copilot/tools/handlers/resources.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/resources.test.ts @@ -4,10 +4,12 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { listAllWorkspaceFilesMock, readWorkspaceFileMetadataMock } = vi.hoisted(() => ({ - listAllWorkspaceFilesMock: vi.fn(), - readWorkspaceFileMetadataMock: vi.fn(), -})) +const { listAllWorkspaceFilesMock, readKnowledgeBaseMock, readWorkspaceFileMetadataMock } = + vi.hoisted(() => ({ + listAllWorkspaceFilesMock: vi.fn(), + readKnowledgeBaseMock: vi.fn(), + readWorkspaceFileMetadataMock: vi.fn(), + })) vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => ({ findWorkspaceFileRecord: ( @@ -37,8 +39,11 @@ vi.mock('@/lib/table/service', () => ({ getTableById: vi.fn(), })) -vi.mock('@/lib/knowledge/service', () => ({ - getKnowledgeBaseById: vi.fn(), +vi.mock('@/lib/knowledge/application/knowledge-bases', () => ({ + readKnowledgeBase: { + operation: { id: 'knowledge.read' }, + execute: readKnowledgeBaseMock, + }, })) vi.mock('@/lib/logs/service', () => ({ @@ -136,4 +141,55 @@ describe('executeOpenResource', () => { ], }) }) + + it('opens a knowledge base through trusted application delegation', async () => { + readKnowledgeBaseMock.mockResolvedValue({ + knowledgeBase: { id: 'kb-1', name: 'Product Docs', workspaceId: 'workspace-1' }, + folderPath: '/', + }) + + const result = await executeOpenResource( + { resources: [{ type: 'knowledgebase', id: 'kb-1' }] }, + { + userId: 'user-1', + workflowId: 'workflow-1', + workspaceId: 'workspace-1', + toolCallId: 'tool-1', + copilotToolExecution: true, + } + ) + + expect(readKnowledgeBaseMock).toHaveBeenCalledWith( + expect.objectContaining({ + principal: expect.objectContaining({ + kind: 'delegated', + subjectUserId: 'user-1', + workspaceId: 'workspace-1', + delegationId: 'tool-1', + }), + input: { knowledgeBaseId: 'kb-1', assertedWorkspaceId: 'workspace-1' }, + }) + ) + expect(result).toMatchObject({ + success: true, + resources: [{ type: 'knowledgebase', id: 'kb-1', title: 'Product Docs' }], + }) + }) + + it('propagates knowledge application infrastructure failures', async () => { + readKnowledgeBaseMock.mockRejectedValueOnce(new Error('knowledge database unavailable')) + + await expect( + executeOpenResource( + { resources: [{ type: 'knowledgebase', id: 'kb-1' }] }, + { + userId: 'user-1', + workflowId: 'workflow-1', + workspaceId: 'workspace-1', + toolCallId: 'tool-1', + copilotToolExecution: true, + } + ) + ).rejects.toThrow('knowledge database unavailable') + }) }) diff --git a/apps/sim/lib/copilot/tools/handlers/resources.ts b/apps/sim/lib/copilot/tools/handlers/resources.ts index 6a5e5556d4b..743c67b7784 100644 --- a/apps/sim/lib/copilot/tools/handlers/resources.ts +++ b/apps/sim/lib/copilot/tools/handlers/resources.ts @@ -1,8 +1,10 @@ import { executeCopilotFileUseCase } from '@/lib/copilot/application/execute-file-use-case' +import { executeCopilotKnowledgeUseCase } from '@/lib/copilot/application/execute-knowledge-use-case' import type { ExecutionContext, ToolCallResult } from '@/lib/copilot/request/types' import { type MothershipResource, MothershipResourceType } from '@/lib/copilot/resources/types' import { canonicalWorkspaceFilePath } from '@/lib/copilot/vfs/path-utils' -import { getKnowledgeBaseById } from '@/lib/knowledge/service' +import { asOrchestrationError } from '@/lib/core/orchestration/types' +import { readKnowledgeBase } from '@/lib/knowledge/application/knowledge-bases' import { getLogById } from '@/lib/logs/service' import { getTableById } from '@/lib/table/service' import { @@ -77,10 +79,27 @@ async function resolveResource( } if (resourceType === 'knowledgebase') { if (!item.id) return { error: 'knowledgebase resources require `id`.' } - const kb = await getKnowledgeBaseById(item.id) - if (!kb) return { error: `No knowledge base with id "${item.id}".` } - if (context.workspaceId && kb.workspaceId !== context.workspaceId) - return { error: `Knowledge base not found in the current workspace.` } + if (!context.workspaceId) { + return { error: 'Opening a knowledge base requires workspace context.' } + } + let kb: Awaited>['knowledgeBase'] + try { + const result = await executeCopilotKnowledgeUseCase(context, readKnowledgeBase, { + knowledgeBaseId: item.id, + assertedWorkspaceId: context.workspaceId, + }) + kb = result.knowledgeBase + } catch (error) { + const classified = asOrchestrationError(error) + if ( + classified?.code === 'not_found' || + classified?.code === 'forbidden' || + classified?.code === 'unauthorized' + ) { + return { error: 'Knowledge base not found in the current workspace.' } + } + throw error + } resourceId = kb.id title = kb.name } diff --git a/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts b/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts index 9e3665e809a..4dd79f4fffc 100644 --- a/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts @@ -33,9 +33,10 @@ const mocks = vi.hoisted(() => ({ verifyFolderWorkspace: vi.fn(), listTables: vi.fn(), renameTable: vi.fn(), - getKnowledgeBases: vi.fn(), + listKnowledgeBases: vi.fn(), updateKnowledgeBase: vi.fn(), - checkKnowledgeBaseWriteAccess: vi.fn(), + deleteKnowledgeBase: vi.fn(), + knowledgeBaseDeleted: vi.fn(), })) vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock })) @@ -151,17 +152,28 @@ vi.mock('@/lib/table/service', () => ({ renameTable: mocks.renameTable, })) -vi.mock('@/lib/knowledge/service', () => ({ - getKnowledgeBases: mocks.getKnowledgeBases, - updateKnowledgeBase: mocks.updateKnowledgeBase, +vi.mock('@/lib/knowledge/application/knowledge-bases', () => ({ + listKnowledgeBases: { + operation: { id: 'knowledge.list' }, + execute: mocks.listKnowledgeBases, + }, + updateKnowledgeBaseOperation: { + operation: { id: 'knowledge.update' }, + execute: mocks.updateKnowledgeBase, + }, + deleteKnowledgeBaseOperation: { + operation: { id: 'knowledge.delete' }, + execute: mocks.deleteKnowledgeBase, + }, })) -vi.mock('@/app/api/knowledge/utils', () => ({ - checkKnowledgeBaseWriteAccess: mocks.checkKnowledgeBaseWriteAccess, +vi.mock('@/lib/core/telemetry', () => ({ + PlatformEvents: { knowledgeBaseDeleted: mocks.knowledgeBaseDeleted }, })) import type { ExecutionContext } from '@/lib/copilot/request/types' -import { executeVfsCp, executeVfsMkdir, executeVfsMv } from './vfs-mutate' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { executeVfsCp, executeVfsMkdir, executeVfsMv, executeVfsRm } from './vfs-mutate' const context = { userId: 'user-1', @@ -594,25 +606,69 @@ describe('vfs mv/cp', () => { expect(result.error).toContain('cannot be copied') }) - it('renames a knowledge base after a write-access check', async () => { - mocks.getKnowledgeBases.mockResolvedValue([{ id: 'kb-1', name: 'Docs' }]) - mocks.checkKnowledgeBaseWriteAccess.mockResolvedValue({ hasAccess: true }) - mocks.updateKnowledgeBase.mockResolvedValue({ id: 'kb-1', name: 'Product Docs' }) + it('renames a knowledge base through trusted application operations', async () => { + mocks.listKnowledgeBases.mockResolvedValue({ + knowledgeBases: [{ knowledgeBase: { id: 'kb-1', name: 'Docs' }, folderPath: '/' }], + }) + mocks.updateKnowledgeBase.mockResolvedValue({ + knowledgeBase: { id: 'kb-1', name: 'Product Docs' }, + folderPath: '/', + }) const result = await executeVfsMv( { sources: ['knowledgebases/Docs'], destination: 'knowledgebases/Product Docs' }, context ) - expect(mocks.checkKnowledgeBaseWriteAccess).toHaveBeenCalledWith('kb-1', 'user-1') expect(mocks.updateKnowledgeBase).toHaveBeenCalledWith( - 'kb-1', - { name: 'Product Docs' }, - expect.any(String) + expect.objectContaining({ + principal: expect.objectContaining({ + kind: 'delegated', + subjectUserId: 'user-1', + workspaceId: 'ws-1', + delegationId: 'tool-call-1', + }), + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: 'ws-1', + name: 'Product Docs', + source: 'agent', + }, + }) ) expect(result.success).toBe(true) }) + it('propagates knowledge application infrastructure failures', async () => { + mocks.listKnowledgeBases.mockRejectedValueOnce(new Error('knowledge database unavailable')) + + await expect( + executeVfsMv( + { sources: ['knowledgebases/Docs'], destination: 'knowledgebases/Product Docs' }, + context + ) + ).rejects.toThrow('knowledge database unavailable') + }) + + it('preserves an actionable knowledge rename conflict', async () => { + mocks.listKnowledgeBases.mockResolvedValue({ + knowledgeBases: [{ knowledgeBase: { id: 'kb-1', name: 'Docs' }, folderPath: '/' }], + }) + mocks.updateKnowledgeBase.mockRejectedValue( + new OrchestrationError('conflict', 'A knowledge base named Product Docs already exists') + ) + + const result = await executeVfsMv( + { sources: ['knowledgebases/Docs'], destination: 'knowledgebases/Product Docs' }, + context + ) + + expect(result).toMatchObject({ + success: false, + error: 'A knowledge base named Product Docs already exists', + }) + }) + it('rejects the reserved knowledgebases/connectors name', async () => { const result = await executeVfsMv( { sources: ['knowledgebases/Docs'], destination: 'knowledgebases/connectors' }, @@ -621,5 +677,54 @@ describe('vfs mv/cp', () => { expect(result.success).toBe(false) expect(result.error).toContain('reserved') }) + + it('deletes a knowledge base through the trusted application operation', async () => { + mocks.listKnowledgeBases.mockResolvedValue({ + knowledgeBases: [{ knowledgeBase: { id: 'kb-1', name: 'Docs' }, folderPath: '/' }], + }) + mocks.deleteKnowledgeBase.mockResolvedValue({ id: 'kb-1', name: 'Docs' }) + + const result = await executeVfsRm({ paths: ['knowledgebases/Docs'] }, context) + + expect(result).toMatchObject({ + success: true, + output: { results: [{ from: 'knowledgebases/Docs', id: 'kb-1' }] }, + }) + expect(mocks.deleteKnowledgeBase).toHaveBeenCalledWith( + expect.objectContaining({ + principal: expect.objectContaining({ delegationId: 'tool-call-1' }), + input: { + knowledgeBaseId: 'kb-1', + assertedWorkspaceId: 'ws-1', + source: 'agent', + }, + }) + ) + expect(mocks.knowledgeBaseDeleted).toHaveBeenCalledWith({ knowledgeBaseId: 'kb-1' }) + }) + + it('preserves an actionable knowledge delete failure', async () => { + mocks.listKnowledgeBases.mockResolvedValue({ + knowledgeBases: [{ knowledgeBase: { id: 'kb-1', name: 'Docs' }, folderPath: '/' }], + }) + mocks.deleteKnowledgeBase.mockRejectedValue( + new OrchestrationError('not_found', 'Knowledge base no longer exists') + ) + + const result = await executeVfsRm({ paths: ['knowledgebases/Docs'] }, context) + + expect(result).toMatchObject({ + success: false, + error: 'Knowledge base no longer exists', + output: { + results: [ + expect.objectContaining({ + from: 'knowledgebases/Docs', + error: 'Knowledge base no longer exists', + }), + ], + }, + }) + }) }) }) diff --git a/apps/sim/lib/copilot/tools/handlers/vfs-mutate.ts b/apps/sim/lib/copilot/tools/handlers/vfs-mutate.ts index d0a4779348c..961545fb299 100644 --- a/apps/sim/lib/copilot/tools/handlers/vfs-mutate.ts +++ b/apps/sim/lib/copilot/tools/handlers/vfs-mutate.ts @@ -7,6 +7,11 @@ import { executeCopilotFileUseCase, resolveCopilotWorkspaceFileReference, } from '@/lib/copilot/application/execute-file-use-case' +import { + executeCopilotKnowledgeUseCase, + messageForCopilotKnowledgeError, + resolveCopilotKnowledgePrincipal, +} from '@/lib/copilot/application/execute-knowledge-use-case' import { messageForCopilotFileError } from '@/lib/copilot/auth/file-delegation' import type { ExecutionContext, ToolCallResult } from '@/lib/copilot/request/types' import { ensureWorkflowAccess, ensureWorkspaceAccess } from '@/lib/copilot/tools/handlers/access' @@ -22,13 +27,14 @@ import { encodeVfsPathSegments, } from '@/lib/copilot/vfs/path-utils' import { asOrchestrationError } from '@/lib/core/orchestration/types' +import { PlatformEvents } from '@/lib/core/telemetry' import { generateRequestId } from '@/lib/core/utils/request' import { createFolder, deleteFolder, updateFolder } from '@/lib/folders/orchestration' import { - deleteKnowledgeBase, - getKnowledgeBases, - updateKnowledgeBase, -} from '@/lib/knowledge/service' + deleteKnowledgeBaseOperation, + listKnowledgeBases, + updateKnowledgeBaseOperation, +} from '@/lib/knowledge/application/knowledge-bases' import { performDeleteTable, performRenameTable } from '@/lib/table/orchestration' import { listTables } from '@/lib/table/service' import { findWorkspaceFileFolderIdByPath } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager' @@ -42,7 +48,6 @@ import { moveWorkspaceFileItemsOperation } from '@/lib/workspace-files/applicati import { fileOperations } from '@/lib/workspace-files/application/operations' import { renameWorkspaceFile } from '@/lib/workspace-files/application/rename-workspace-file' import { updateWorkspaceFileFolderOperation } from '@/lib/workspace-files/application/workspace-file-folders' -import { checkKnowledgeBaseWriteAccess } from '@/app/api/knowledge/utils' const logger = createLogger('VfsMutateTools') @@ -78,6 +83,21 @@ interface VfsMutateOutcome { error?: string } +class KnowledgeVfsInfrastructureError extends Error { + constructor(readonly infrastructureCause: unknown) { + super('Knowledge VFS infrastructure failure') + this.name = 'KnowledgeVfsInfrastructureError' + } +} + +function messageForKnowledgeVfsError(error: unknown, forbiddenMessage: string): string { + const classified = asOrchestrationError(error) + if (!classified || classified.code === 'internal') { + throw new KnowledgeVfsInfrastructureError(error) + } + return classified.code === 'forbidden' ? forbiddenMessage : messageForCopilotKnowledgeError(error) +} + /** Top-level VFS segment of a raw (possibly encoded) path. */ function topLevelSegment(path: string): string { return path.trim().replace(/^\/+/, '').split('/')[0] ?? '' @@ -239,6 +259,9 @@ async function executeVfsMutate( } const workspaceId = requireCopilotWorkspace(context) + if (topLevelSegment(sources[0]) === 'knowledgebases') { + resolveCopilotKnowledgePrincipal(context) + } await ensureWorkspaceAccess(workspaceId, context.userId, 'write') assertMutationNotAborted(context) @@ -274,6 +297,9 @@ async function executeVfsMutate( return await renameFlatResource(verb, category, sources, destination, context, workspaceId) } } catch (error) { + if (error instanceof KnowledgeVfsInfrastructureError) { + throw error.infrastructureCause + } return { success: false, error: context.abortSignal?.aborted @@ -837,20 +863,41 @@ async function renameFlatResource( if (newName.toLowerCase() === 'connectors') { return { success: false, error: '"knowledgebases/connectors" is a reserved path.' } } - const kbs = await getKnowledgeBases(context.userId, workspaceId) - const match = kbs.find((kb) => normalizeVfsSegment(kb.name) === canonicalSource) + let knowledgeBases: Awaited>['knowledgeBases'] + try { + const result = await executeCopilotKnowledgeUseCase(context, listKnowledgeBases, { + workspaceId, + }) + knowledgeBases = result.knowledgeBases + } catch (error) { + return { + success: false, + error: messageForKnowledgeVfsError(error, 'Write access required to rename knowledge bases'), + } + } + const match = knowledgeBases + .map(({ knowledgeBase }) => knowledgeBase) + .find((kb) => normalizeVfsSegment(kb.name) === canonicalSource) if (!match) { return { success: false, error: `Knowledge base not found at ${sources[0]}` } } - const access = await checkKnowledgeBaseWriteAccess(match.id, context.userId) - if (!access.hasAccess) { + assertMutationNotAborted(context) + try { + await executeCopilotKnowledgeUseCase(context, updateKnowledgeBaseOperation, { + knowledgeBaseId: match.id, + assertedWorkspaceId: workspaceId, + name: newName, + source: 'agent', + }) + } catch (error) { return { success: false, - error: `Write access required to rename knowledge base "${match.name}"`, + error: messageForKnowledgeVfsError( + error, + `Write access required to rename knowledge base "${match.name}"` + ), } } - assertMutationNotAborted(context) - await updateKnowledgeBase(match.id, { name: newName }, generateRequestId()) logger.info('Renamed knowledge base via mv', { knowledgeBaseId: match.id, workspaceId }) return buildResult(verb, [ { from: sources[0], to: `knowledgebases/${normalizeVfsSegment(newName)}`, kind, id: match.id }, @@ -877,6 +924,9 @@ export async function executeVfsRm( } const workspaceId = requireCopilotWorkspace(context) + if (paths.some((path) => topLevelSegment(path) === 'knowledgebases')) { + resolveCopilotKnowledgePrincipal(context) + } await ensureWorkspaceAccess(workspaceId, context.userId, 'write') assertMutationNotAborted(context) @@ -897,6 +947,7 @@ export async function executeVfsRm( await removeOne(classified.category, path, context, workspaceId, getWorkflowIndex) ) } catch (error) { + if (error instanceof KnowledgeVfsInfrastructureError) throw error outcomes.push({ from: path, kind: defaultKindFor(path), @@ -910,6 +961,9 @@ export async function executeVfsRm( return buildResult('rm', outcomes) } catch (error) { + if (error instanceof KnowledgeVfsInfrastructureError) { + throw error.infrastructureCause + } return { success: false, error: context.abortSignal?.aborted @@ -1139,23 +1193,43 @@ async function removeKnowledgeBasePath( error: '"knowledgebases/connectors" is a reserved path, not a knowledge base.', } } - const match = (await getKnowledgeBases(context.userId, workspaceId)).find( - (kb) => normalizeVfsSegment(kb.name) === canonical - ) + let knowledgeBases: Awaited>['knowledgeBases'] + try { + const result = await executeCopilotKnowledgeUseCase(context, listKnowledgeBases, { + workspaceId, + }) + knowledgeBases = result.knowledgeBases + } catch (error) { + return { + from: path, + kind: 'knowledge_base', + error: messageForKnowledgeVfsError(error, 'Write access required to delete knowledge bases'), + } + } + const match = knowledgeBases + .map(({ knowledgeBase }) => knowledgeBase) + .find((kb) => normalizeVfsSegment(kb.name) === canonical) if (!match) return { from: path, kind: 'knowledge_base', error: `Knowledge base not found at ${path}` } - const access = await checkKnowledgeBaseWriteAccess(match.id, context.userId) - if (!access.hasAccess) { + try { + await executeCopilotKnowledgeUseCase(context, deleteKnowledgeBaseOperation, { + knowledgeBaseId: match.id, + assertedWorkspaceId: workspaceId, + source: 'agent', + }) + } catch (error) { return { from: path, kind: 'knowledge_base', id: match.id, - error: `Write access required to delete knowledge base "${match.name}"`, + error: messageForKnowledgeVfsError( + error, + `Write access required to delete knowledge base "${match.name}"` + ), } } - - await deleteKnowledgeBase(match.id, generateRequestId()) + PlatformEvents.knowledgeBaseDeleted({ knowledgeBaseId: match.id }) logger.info('Deleted knowledge base via rm', { knowledgeBaseId: match.id, workspaceId }) return { from: path, kind: 'knowledge_base', id: match.id } } diff --git a/apps/sim/lib/copilot/tools/handlers/vfs.test.ts b/apps/sim/lib/copilot/tools/handlers/vfs.test.ts index 3d2d82ed9bd..21284cb77ed 100644 --- a/apps/sim/lib/copilot/tools/handlers/vfs.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/vfs.test.ts @@ -75,7 +75,13 @@ function makeVfs() { } } -const GREP_CTX = { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } +const GREP_CTX = { + userId: 'user-1', + workflowId: 'wf-1', + workspaceId: 'ws-1', + toolCallId: 'tool-1', + copilotToolExecution: true, +} const GREP_CTX_CHAT = { ...GREP_CTX, chatId: 'chat-1' } describe('vfs handlers oversize policy', () => { @@ -89,10 +95,7 @@ describe('vfs handlers oversize policy', () => { vfs.grep.mockReturnValue([{ path: 'files/a.txt', line: 1, content: OVERSIZED_INLINE_CONTENT }]) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsGrep( - { pattern: 'foo', output_mode: 'content' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsGrep({ pattern: 'foo', output_mode: 'content' }, GREP_CTX) expect(result.success).toBe(false) expect(result.error).toContain('more specific pattern') @@ -105,10 +108,7 @@ describe('vfs handlers oversize policy', () => { vfs.read.mockReturnValue({ content: OVERSIZED_INLINE_CONTENT, totalLines: 1 }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'workflows/My Workflow/state.json' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'workflows/My Workflow/state.json' }, GREP_CTX) expect(result.success).toBe(false) expect(result.error).toContain('Use grep') @@ -124,10 +124,7 @@ describe('vfs handlers oversize policy', () => { }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'files/big.txt/content' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'files/big.txt/content' }, GREP_CTX) expect(result.success).toBe(false) expect(result.error).toContain('File too large to display inline') @@ -147,10 +144,7 @@ describe('vfs handlers oversize policy', () => { }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'files/chess.png/content' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'files/chess.png/content' }, GREP_CTX) expect(result.success).toBe(true) expect((result.output as { attachment?: { type: string } })?.attachment?.type).toBe('image') @@ -170,10 +164,7 @@ describe('vfs handlers oversize policy', () => { }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'files/reports/report.pdf/compiled' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'files/reports/report.pdf/compiled' }, GREP_CTX) expect(result.success).toBe(true) expect((result.output as { attachment?: { type: string } })?.attachment?.type).toBe('file') @@ -187,10 +178,7 @@ describe('vfs handlers oversize policy', () => { }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'files/huge.png/content' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'files/huge.png/content' }, GREP_CTX) expect(result.success).toBe(false) expect(result.error).toContain('too large') @@ -204,10 +192,7 @@ describe('vfs handlers oversize policy', () => { }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'files/report.csv' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'files/report.csv' }, GREP_CTX) expect(result.success).toBe(true) expect(vfs.readFileContent).not.toHaveBeenCalled() @@ -229,9 +214,18 @@ describe('vfs handlers oversize policy', () => { ) expect(result.success).toBe(true) - expect(getOrMaterializeVFS).toHaveBeenCalledWith('ws-1', 'user-1', { - secretMountPolicy, - }) + expect(getOrMaterializeVFS).toHaveBeenCalledWith( + 'ws-1', + 'user-1', + expect.objectContaining({ + secretMountPolicy, + knowledgePrincipal: expect.objectContaining({ + kind: 'delegated', + delegationId: 'tool-1', + workspaceId: 'ws-1', + }), + }) + ) }) it('uses dynamic file reads for canonical style paths', async () => { @@ -242,10 +236,7 @@ describe('vfs handlers oversize policy', () => { }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'files/reports/brief.docx/style' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'files/reports/brief.docx/style' }, GREP_CTX) expect(result.success).toBe(true) expect(vfs.readFileContent).toHaveBeenCalledWith('files/reports/brief.docx/style') @@ -260,10 +251,7 @@ describe('vfs handlers oversize policy', () => { }) getOrMaterializeVFS.mockResolvedValue(vfs) - const result = await executeVfsRead( - { path: 'files/reports/brief.pdf/compiled' }, - { userId: 'user-1', workflowId: 'wf-1', workspaceId: 'ws-1' } - ) + const result = await executeVfsRead({ path: 'files/reports/brief.pdf/compiled' }, GREP_CTX) expect(result.success).toBe(true) expect(vfs.readFileContent).toHaveBeenCalledWith('files/reports/brief.pdf/compiled') diff --git a/apps/sim/lib/copilot/tools/handlers/vfs.ts b/apps/sim/lib/copilot/tools/handlers/vfs.ts index 3aade59a519..b3bb30cbe6a 100644 --- a/apps/sim/lib/copilot/tools/handlers/vfs.ts +++ b/apps/sim/lib/copilot/tools/handlers/vfs.ts @@ -1,5 +1,6 @@ import { createLogger } from '@sim/logger' import { getErrorMessage, toError } from '@sim/utils/errors' +import { resolveCopilotKnowledgePrincipal } from '@/lib/copilot/application/execute-knowledge-use-case' import { resolveCopilotFilePrincipal } from '@/lib/copilot/auth/file-delegation' import { getBlockVisibilityForCopilot } from '@/lib/copilot/block-visibility' import { TOOL_RESULT_MAX_INLINE_CHARS } from '@/lib/copilot/constants' @@ -30,6 +31,7 @@ const logger = createLogger('VfsTools') async function getGatedVFS(context: ExecutionContext) { const workspaceId = context.workspaceId if (!workspaceId) throw new Error('No workspace context available') + const knowledgePrincipal = resolveCopilotKnowledgePrincipal(context) const vis = await getBlockVisibilityForCopilot(context.userId, workspaceId) const filePrincipal = context.copilotToolExecution && context.toolCallId @@ -39,6 +41,7 @@ async function getGatedVFS(context: ExecutionContext) { getOrMaterializeVFS(workspaceId, context.userId, { secretMountPolicy: context.secretMountPolicy, filePrincipal, + knowledgePrincipal, }) ) } diff --git a/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.test.ts b/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.test.ts index fc3b1a54ae1..7e850b80dd5 100644 --- a/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.test.ts +++ b/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.test.ts @@ -1,69 +1,68 @@ /** * @vitest-environment node */ -import { knowledgeConnector } from '@sim/db/schema' -import { loggerMock, queueTableRows, resetDbChainMock } from '@sim/testing' -import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' +import { beforeEach, describe, expect, it, vi } from 'vitest' const { - mockAssertBillingAttributionSnapshot, - mockCheckKnowledgeBaseWriteAccess, - mockGetKnowledgeBaseById, + mockCaptureServerEvent, + mockCreateKnowledgeBase, + mockDeleteKnowledgeBase, + mockDeleteKnowledgeDocument, mockGetBoundWorkspaceFileSecretProvenance, - mockImportKnowledgeSearchResultSecretProvenance, - mockPerformCreateKnowledgeConnector, - mockPerformDeleteKnowledgeBase, - mockPerformDeleteKnowledgeConnector, - mockPerformSyncKnowledgeConnector, + mockKnowledgeBaseCreated, + mockKnowledgeBaseDeleted, + mockKnowledgeBaseDocumentsUploaded, + mockReadKnowledgeBase, + mockResolveWorkspaceFileReference, + mockSearchKnowledge, + mockUpdateKnowledgeBase, + mockUploadKnowledgeDocument, } = vi.hoisted(() => ({ - mockAssertBillingAttributionSnapshot: vi.fn(), - mockCheckKnowledgeBaseWriteAccess: vi.fn(), - mockGetKnowledgeBaseById: vi.fn(), + mockCaptureServerEvent: vi.fn(), + mockCreateKnowledgeBase: vi.fn(), + mockDeleteKnowledgeBase: vi.fn(), + mockDeleteKnowledgeDocument: vi.fn(), mockGetBoundWorkspaceFileSecretProvenance: vi.fn(), - mockImportKnowledgeSearchResultSecretProvenance: vi.fn(), - mockPerformCreateKnowledgeConnector: vi.fn(), - mockPerformDeleteKnowledgeBase: vi.fn(), - mockPerformDeleteKnowledgeConnector: vi.fn(), - mockPerformSyncKnowledgeConnector: vi.fn(), + mockKnowledgeBaseCreated: vi.fn(), + mockKnowledgeBaseDeleted: vi.fn(), + mockKnowledgeBaseDocumentsUploaded: vi.fn(), + mockReadKnowledgeBase: vi.fn(), + mockResolveWorkspaceFileReference: vi.fn(), + mockSearchKnowledge: vi.fn(), + mockUpdateKnowledgeBase: vi.fn(), + mockUploadKnowledgeDocument: vi.fn(), })) -vi.mock('@/lib/billing/calculations/usage-monitor', () => ({ - checkActorUsageLimits: vi.fn(), -})) -vi.mock('@/lib/billing/core/billing-attribution', () => ({ - assertBillingAttributionSnapshot: mockAssertBillingAttributionSnapshot, - checkAttributedUsageLimits: vi.fn(), -})) vi.mock('@/lib/copilot/generated/tool-catalog-v1', () => ({ KnowledgeBase: { id: 'knowledge_base' }, })) -vi.mock('@/lib/copilot/tools/server/base-tool', () => ({ - assertServerToolNotAborted: vi.fn(), +vi.mock('@/lib/core/telemetry', () => ({ + PlatformEvents: { + knowledgeBaseCreated: mockKnowledgeBaseCreated, + knowledgeBaseDeleted: mockKnowledgeBaseDeleted, + knowledgeBaseDocumentsUploaded: mockKnowledgeBaseDocumentsUploaded, + }, +})) +vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: mockCaptureServerEvent })) +vi.mock('@/lib/knowledge/application/knowledge-bases', () => ({ + createKnowledgeBase: { execute: mockCreateKnowledgeBase }, + deleteKnowledgeBaseOperation: { execute: mockDeleteKnowledgeBase }, + readKnowledgeBase: { execute: mockReadKnowledgeBase }, + updateKnowledgeBaseOperation: { execute: mockUpdateKnowledgeBase }, +})) +vi.mock('@/lib/knowledge/application/documents', () => ({ + deleteKnowledgeDocument: { execute: mockDeleteKnowledgeDocument }, + uploadKnowledgeDocument: { execute: mockUploadKnowledgeDocument }, })) -vi.mock('@/lib/knowledge/embeddings', () => ({ - generateSearchEmbedding: vi.fn(), - recordSearchEmbeddingUsage: vi.fn(), +vi.mock('@/lib/knowledge/application/search', () => ({ + searchKnowledge: { execute: mockSearchKnowledge }, })) vi.mock('@/lib/knowledge/orchestration', () => ({ - performCreateKnowledgeBase: vi.fn(), - performDeleteKnowledgeBase: mockPerformDeleteKnowledgeBase, - performCreateKnowledgeConnector: mockPerformCreateKnowledgeConnector, - performDeleteKnowledgeConnector: mockPerformDeleteKnowledgeConnector, - performDeleteKnowledgeDocument: vi.fn(), - performSyncKnowledgeConnector: mockPerformSyncKnowledgeConnector, - performUpdateKnowledgeBase: vi.fn(), + performCreateKnowledgeConnector: vi.fn(), + performDeleteKnowledgeConnector: vi.fn(), + performSyncKnowledgeConnector: vi.fn(), performUpdateKnowledgeConnector: vi.fn(), performUpdateKnowledgeDocument: vi.fn(), - performUploadKnowledgeDocument: vi.fn(), -})) -vi.mock('@/lib/knowledge/service', () => ({ - getKnowledgeBaseById: mockGetKnowledgeBaseById, -})) -vi.mock('@/lib/knowledge/secret-provenance', () => ({ - importKnowledgeSearchResultSecretProvenance: mockImportKnowledgeSearchResultSecretProvenance, -})) -vi.mock('@/lib/knowledge/documents/service', () => ({ - createSingleDocument: vi.fn(), })) vi.mock('@/lib/knowledge/tags/service', () => ({ createTagDefinition: vi.fn(), @@ -74,438 +73,388 @@ vi.mock('@/lib/knowledge/tags/service', () => ({ getTagUsageStats: vi.fn(), updateTagDefinition: vi.fn(), })) -vi.mock('@/lib/uploads', () => ({ StorageService: {} })) +vi.mock('@/lib/uploads', () => ({ + StorageService: { generatePresignedDownloadUrl: vi.fn().mockResolvedValue('https://file.test') }, +})) vi.mock('@/lib/workspace-files/application/resolve-workspace-file-reference', () => ({ - resolveWorkspaceFileReference: vi.fn(), + resolveWorkspaceFileReference: mockResolveWorkspaceFileReference, })) -vi.mock('@/app/api/auth/oauth/utils', () => ({ getCredential: vi.fn() })) vi.mock('@/lib/uploads/contexts/workspace/workspace-file-secret-provenance', () => ({ getBoundWorkspaceFileSecretProvenance: mockGetBoundWorkspaceFileSecretProvenance, })) -vi.mock('@/lib/knowledge/search/queries', () => ({ - executeKnowledgeSearch: vi.fn(), -})) +vi.mock('@/app/api/auth/oauth/utils', () => ({ getCredential: vi.fn() })) vi.mock('@/app/api/knowledge/utils', () => ({ checkDocumentWriteAccess: vi.fn(), checkKnowledgeBaseAccess: vi.fn(), - checkKnowledgeBaseWriteAccess: mockCheckKnowledgeBaseWriteAccess, + checkKnowledgeBaseWriteAccess: vi.fn(), })) -import { checkAttributedUsageLimits } from '@/lib/billing/core/billing-attribution' -import { projectToolResultForCopilot } from '@/lib/copilot/request/tools/resolved-secret-result' +import type { ServerToolContext } from '@/lib/copilot/tools/server/base-tool' import { knowledgeBaseServerTool } from '@/lib/copilot/tools/server/knowledge/knowledge-base' -import { createSingleDocument } from '@/lib/knowledge/documents/service' -import { generateSearchEmbedding, recordSearchEmbeddingUsage } from '@/lib/knowledge/embeddings' -import { executeKnowledgeSearch } from '@/lib/knowledge/search/queries' -import { getKnowledgeBaseById } from '@/lib/knowledge/service' -import { resolveWorkspaceFileReference } from '@/lib/workspace-files/application/resolve-workspace-file-reference' -import { checkKnowledgeBaseAccess } from '@/app/api/knowledge/utils' +import { OrchestrationError } from '@/lib/core/orchestration/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' -const knowledgeLoggerIndex = loggerMock.createLogger.mock.calls.findIndex( - ([name]) => name === 'KnowledgeBaseServerTool' -) -const knowledgeLogger = loggerMock.createLogger.mock.results[knowledgeLoggerIndex]?.value - -const BILLING_ATTRIBUTION = { - actorUserId: 'external-admin', +const KNOWLEDGE_BASE = { + id: 'knowledge-base-1', + name: 'Private KB', + description: 'Private documentation', workspaceId: 'workspace-paid', - organizationId: 'organization-paid', - billedAccountUserId: 'workspace-owner', - billingEntity: { type: 'organization' as const, id: 'organization-paid' }, - billingPeriod: { - start: '2026-07-01T00:00:00.000Z', - end: '2026-08-01T00:00:00.000Z', - }, - payerSubscription: null, + docCount: 2, + tokenCount: 42, + embeddingModel: 'text-embedding-3-small', + chunkingConfig: { maxSize: 1024, minSize: 100, overlap: 200 }, + createdAt: new Date('2026-08-01T00:00:00.000Z'), + updatedAt: new Date('2026-08-02T00:00:00.000Z'), } const CONTEXT = { userId: 'external-admin', workspaceId: 'workspace-paid', - billingAttribution: BILLING_ATTRIBUTION, -} - -describe('knowledge base connector Copilot operations', () => { - afterAll(() => { - resetDbChainMock() + chatId: 'chat-1', + executionId: 'execution-1', + toolCallId: 'tool-call-1', + copilotToolExecution: true, +} satisfies ServerToolContext + +function expectDelegatedPrincipal(call: unknown): void { + expect(call).toMatchObject({ + principal: { + kind: 'delegated', + serviceId: 'copilot', + subjectUserId: 'external-admin', + workspaceId: 'workspace-paid', + delegationId: 'tool-call-1', + audience: 'sim:knowledge', + resourceScope: { chatId: 'chat-1', executionId: 'execution-1' }, + }, }) +} +describe('knowledge_base trusted application delegation', () => { beforeEach(() => { vi.clearAllMocks() - resetDbChainMock() - queueTableRows(knowledgeConnector, [{ knowledgeBaseId: 'knowledge-base-1' }]) - mockAssertBillingAttributionSnapshot.mockReturnValue(BILLING_ATTRIBUTION) - mockCheckKnowledgeBaseWriteAccess.mockResolvedValue({ - hasAccess: true, - knowledgeBase: { - id: 'knowledge-base-1', - workspaceId: 'workspace-paid', - name: 'Paid KB', - }, - }) - mockPerformCreateKnowledgeConnector.mockResolvedValue({ - success: true, - connector: { id: 'connector-1', connectorType: 'notion', status: 'active' }, - }) - mockPerformSyncKnowledgeConnector.mockResolvedValue({ success: true }) - mockPerformDeleteKnowledgeConnector.mockResolvedValue({ - success: true, - documentsDeleted: 0, - documentsKept: 3, + mockReadKnowledgeBase.mockResolvedValue({ knowledgeBase: KNOWLEDGE_BASE, folderPath: '/' }) + mockCreateKnowledgeBase.mockResolvedValue({ knowledgeBase: KNOWLEDGE_BASE, folderPath: '/' }) + mockUpdateKnowledgeBase.mockResolvedValue({ knowledgeBase: KNOWLEDGE_BASE, folderPath: '/' }) + mockDeleteKnowledgeBase.mockResolvedValue({ id: KNOWLEDGE_BASE.id, name: KNOWLEDGE_BASE.name }) + mockSearchKnowledge.mockResolvedValue({ + results: [], + query: 'query', + knowledgeBaseIds: [KNOWLEDGE_BASE.id], + topK: 5, + totalResults: 0, }) + mockDeleteKnowledgeDocument.mockResolvedValue({ id: 'document-1', filename: 'doc.pdf' }) }) it.each([ - { - operation: 'add_connector', - params: { - operation: 'add_connector', - args: { - knowledgeBaseId: 'knowledge-base-1', - connectorType: 'notion', - apiKey: 'api-key', - }, - }, - perform: mockPerformCreateKnowledgeConnector, - }, - { - operation: 'sync_connector', - params: { operation: 'sync_connector', args: { connectorId: 'connector-1' } }, - perform: mockPerformSyncKnowledgeConnector, - }, - ])('forwards immutable billing attribution for $operation', async ({ params, perform }) => { - const result = await knowledgeBaseServerTool.execute(params, CONTEXT) - - expect(result.success).toBe(true) - // The operation runs in-process now. The payer travels as a value on the - // orchestration call rather than as a serialized header on an internal - // HTTP self-call back into this same process. - const call = perform.mock.calls[0][0] - expect(await call.resolveBillingAttribution()).toEqual(BILLING_ATTRIBUTION) - expect(call.source).toBe('agent') - expect(mockAssertBillingAttributionSnapshot).toHaveBeenCalledWith(BILLING_ATTRIBUTION) + [{ ...CONTEXT, copilotToolExecution: false }, 'trusted Copilot execution context'], + [{ ...CONTEXT, workspaceId: undefined }, 'workspace ID'], + [{ ...CONTEXT, toolCallId: undefined }, 'tool call ID'], + [{ ...CONTEXT, userId: '' }, 'authenticated user ID'], + ])('rejects incomplete server-authored context', async (context, message) => { + await expect( + knowledgeBaseServerTool.execute( + { operation: 'get', args: { knowledgeBaseId: KNOWLEDGE_BASE.id } }, + context + ) + ).rejects.toThrow(message) + expect(mockReadKnowledgeBase).not.toHaveBeenCalled() }) - it('reports a failed knowledge base delete as failed, not as missing', async () => { - mockGetKnowledgeBaseById.mockResolvedValue({ - id: 'knowledge-base-1', - name: 'Paid KB', - workspaceId: 'workspace-paid', - }) - mockPerformDeleteKnowledgeBase.mockResolvedValue({ - success: false, - error: 'Knowledge base is locked', - errorCode: 'conflict', - }) - + it('creates in the trusted workspace and ignores a model workspace field', async () => { const result = await knowledgeBaseServerTool.execute( - { operation: 'delete', args: { knowledgeBaseId: 'knowledge-base-1' } }, + { + operation: 'create', + args: { name: 'Private KB', workspaceId: 'model-controlled-workspace' }, + }, CONTEXT ) - // A knowledge base that exists but could not be archived is neither deleted - // nor missing — folding it into notFound told the user it was never there. - expect(result.data.notFound).toEqual([]) - expect(result.data.failed).toEqual([ - { id: 'knowledge-base-1', name: 'Paid KB', reason: 'Knowledge base is locked' }, - ]) - expect(result.message).toContain('Knowledge base is locked') - }) - - it('never relays an unclassified fault to the agent verbatim', async () => { - mockGetKnowledgeBaseById.mockResolvedValue({ - id: 'knowledge-base-1', - name: 'Paid KB', + expect(result.success).toBe(true) + const call = mockCreateKnowledgeBase.mock.calls[0][0] + expectDelegatedPrincipal(call) + expect(call.input).toMatchObject({ workspaceId: 'workspace-paid', + name: 'Private KB', + source: 'agent', }) - mockPerformDeleteKnowledgeBase.mockResolvedValue({ - success: false, - error: 'select "id" from "knowledge_base" — connection terminated', - errorCode: 'internal', + expect(mockKnowledgeBaseCreated).toHaveBeenCalledWith({ + knowledgeBaseId: KNOWLEDGE_BASE.id, + name: KNOWLEDGE_BASE.name, + workspaceId: 'workspace-paid', }) - - const result = await knowledgeBaseServerTool.execute( - { operation: 'delete', args: { knowledgeBaseId: 'knowledge-base-1' } }, - CONTEXT + expect(mockCaptureServerEvent).toHaveBeenCalledWith( + 'external-admin', + 'knowledge_base_created', + expect.objectContaining({ workspace_id: 'workspace-paid' }), + expect.any(Object) ) - - expect(result.data.failed[0].reason).toBe('Failed to delete knowledge base') - expect(result.message).not.toContain('connection terminated') }) - it('reports that a deleted connector kept its documents, because it did', async () => { + it('reads through the canonical application operation', async () => { const result = await knowledgeBaseServerTool.execute( - { operation: 'delete_connector', args: { connectorId: 'connector-1' } }, + { operation: 'get', args: { knowledgeBaseId: KNOWLEDGE_BASE.id } }, CONTEXT ) - // The old wording claimed the documents "have been removed". They never - // were: the tool reached the route over HTTP with no query string, so the - // route's keep-documents default always applied. expect(result.success).toBe(true) - expect(result.message).toContain('3 document(s) were kept') - expect(result.message).not.toContain('removed') - expect(mockPerformDeleteKnowledgeConnector).toHaveBeenCalledWith( - expect.objectContaining({ connectorId: 'connector-1', source: 'agent' }) - ) - }) -}) - -describe('knowledge base query model boundary', () => { - beforeEach(() => { - vi.clearAllMocks() - resetDbChainMock() - vi.mocked(checkKnowledgeBaseAccess).mockResolvedValue({ hasAccess: true }) - vi.mocked(getKnowledgeBaseById).mockResolvedValue({ - id: 'knowledge-base-1', - name: 'Private KB', - workspaceId: 'workspace-paid', - embeddingModel: 'text-embedding-3-small', - } as Awaited>) - vi.mocked(checkAttributedUsageLimits).mockResolvedValue({ isExceeded: false }) - vi.mocked(generateSearchEmbedding).mockResolvedValue({ - embedding: [0.1, 0.2], - isBYOK: false, - }) - vi.mocked(executeKnowledgeSearch).mockResolvedValue([]) - vi.mocked(recordSearchEmbeddingUsage).mockResolvedValue(undefined) - mockImportKnowledgeSearchResultSecretProvenance.mockResolvedValue({ - imported: true, - documentMetadata: {}, + const call = mockReadKnowledgeBase.mock.calls[0][0] + expectDelegatedPrincipal(call) + expect(call.input).toEqual({ + knowledgeBaseId: KNOWLEDGE_BASE.id, + assertedWorkspaceId: 'workspace-paid', }) }) - it('projects the query at embedding, search, and usage boundaries', async () => { + it('projects query secrets before delegating search and passes only the trusted registry', async () => { const registry = new ResolvedSecretTraceRegistry([ { name: 'KB_QUERY', - plaintext: 'private knowledge query', + plaintext: 'private query', encryptedValue: 'encrypted-query', }, ]) - registry.recordResolved('KB_QUERY', 'private knowledge query') + registry.recordResolved('KB_QUERY', 'private query') + mockSearchKnowledge.mockResolvedValueOnce({ + results: [ + { + embeddingId: 'embedding-1', + documentId: 'document-1', + documentName: 'doc.pdf', + sourceUrl: null, + content: 'result', + chunkIndex: 0, + metadata: {}, + similarity: 0.9, + }, + ], + query: '{{KB_QUERY}}', + knowledgeBaseIds: [KNOWLEDGE_BASE.id], + topK: 5, + totalResults: 1, + }) const result = await knowledgeBaseServerTool.execute( { operation: 'query', - args: { - knowledgeBaseId: 'knowledge-base-1', - query: 'private knowledge query', - }, + args: { knowledgeBaseId: KNOWLEDGE_BASE.id, query: 'private query' }, }, - { - userId: 'external-admin', - workspaceId: 'workspace-paid', - toolCallId: 'tool-1', - copilotToolExecution: true, - billingAttribution: BILLING_ATTRIBUTION, - resolvedSecretTraceRegistry: registry, - } + { ...CONTEXT, resolvedSecretTraceRegistry: registry } ) - expect(result.success).toBe(true) - expect(result.data?.query).toBe('private knowledge query') - expect(generateSearchEmbedding).toHaveBeenCalledWith( - '{{KB_QUERY}}', - 'text-embedding-3-small', - 'workspace-paid' - ) - expect(executeKnowledgeSearch).toHaveBeenCalledWith( - expect.objectContaining({ query: '{{KB_QUERY}}' }) - ) - expect(recordSearchEmbeddingUsage).toHaveBeenCalledWith( - expect.objectContaining({ query: '{{KB_QUERY}}' }) - ) - expect(mockImportKnowledgeSearchResultSecretProvenance).toHaveBeenCalledWith({ - registry, - results: [], + expect(result).toMatchObject({ + success: true, + data: { query: 'private query', results: [{ similarity: 0.9 }] }, + }) + const call = mockSearchKnowledge.mock.calls[0][0] + expectDelegatedPrincipal(call) + expect(call.input).toEqual({ + workspaceId: 'workspace-paid', + knowledgeBaseIds: [KNOWLEDGE_BASE.id], + query: '{{KB_QUERY}}', + topK: 5, + resultSecretRegistry: registry, }) - expect(knowledgeLogger).toBeDefined() - expect(JSON.stringify(knowledgeLogger?.info.mock.calls)).not.toContain( - 'private knowledge query' - ) }) - it('imports exact persisted result provenance before the Copilot result is projected', async () => { - const registry = new ResolvedSecretTraceRegistry([ - { - name: 'STORED_TOKEN', - plaintext: 'stored-secret-value', - encryptedValue: 'encrypted-stored-secret', - }, - ]) - const results = [ - { - id: 'embedding-1', - documentId: 'document-1', - content: 'stored-secret-value', - chunkIndex: 0, - distance: 0.1, - }, - ] - vi.mocked(executeKnowledgeSearch).mockResolvedValue(results) - mockImportKnowledgeSearchResultSecretProvenance.mockImplementationOnce( - async ({ registry: resultRegistry }) => { - expect(resultRegistry.recordResolved('STORED_TOKEN', 'stored-secret-value')).toBe(true) - return { imported: true, documentMetadata: {} } - } - ) + it('propagates search infrastructure failures', async () => { + mockSearchKnowledge.mockRejectedValueOnce(new Error('database unavailable')) + + await expect( + knowledgeBaseServerTool.execute( + { operation: 'query', args: { knowledgeBaseId: KNOWLEDGE_BASE.id, query: 'query' } }, + { ...CONTEXT, resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry() } + ) + ).rejects.toThrow('database unavailable') + }) + it('updates through the semantic operation', async () => { const result = await knowledgeBaseServerTool.execute( - { - operation: 'query', - args: { - knowledgeBaseId: 'knowledge-base-1', - query: 'public query', - }, - }, - { - userId: 'external-admin', - workspaceId: 'workspace-paid', - toolCallId: 'tool-1', - copilotToolExecution: true, - billingAttribution: BILLING_ATTRIBUTION, - resolvedSecretTraceRegistry: registry, - } + { operation: 'update', args: { knowledgeBaseId: KNOWLEDGE_BASE.id, name: 'Renamed' } }, + CONTEXT ) expect(result.success).toBe(true) - expect(result.data?.results[0].content).toBe('stored-secret-value') - expect(projectToolResultForCopilot({ success: true, output: result }, registry)).toMatchObject({ + const call = mockUpdateKnowledgeBase.mock.calls[0][0] + expectDelegatedPrincipal(call) + expect(call.input).toMatchObject({ + knowledgeBaseId: KNOWLEDGE_BASE.id, + assertedWorkspaceId: 'workspace-paid', + name: 'Renamed', + source: 'agent', + }) + }) + + it('keeps the unexposed delete compatibility path on the shared delete operation', async () => { + const result = await knowledgeBaseServerTool.execute( + { operation: 'delete', args: { knowledgeBaseId: KNOWLEDGE_BASE.id } }, + CONTEXT + ) + + expect(result).toMatchObject({ success: true, - output: { - data: { results: [{ content: '{{STORED_TOKEN}}' }] }, - }, + data: { deleted: [{ id: KNOWLEDGE_BASE.id, name: KNOWLEDGE_BASE.name }] }, + }) + const call = mockDeleteKnowledgeBase.mock.calls[0][0] + expectDelegatedPrincipal(call) + expect(call.input).toEqual({ + knowledgeBaseId: KNOWLEDGE_BASE.id, + assertedWorkspaceId: 'workspace-paid', + source: 'agent', + }) + expect(mockKnowledgeBaseDeleted).toHaveBeenCalledWith({ + knowledgeBaseId: KNOWLEDGE_BASE.id, }) }) - it('fails closed when persisted result provenance cannot be established', async () => { - const registry = new ResolvedSecretTraceRegistry() - vi.mocked(executeKnowledgeSearch).mockResolvedValue([ - { - id: 'embedding-1', - documentId: 'document-1', - content: 'unclassified persisted content', - chunkIndex: 0, - distance: 0.1, + it('keeps classified delete failures in the batch result', async () => { + mockDeleteKnowledgeBase.mockRejectedValueOnce( + new OrchestrationError('conflict', 'Knowledge base is locked') + ) + + const result = await knowledgeBaseServerTool.execute( + { operation: 'delete', args: { knowledgeBaseId: KNOWLEDGE_BASE.id } }, + CONTEXT + ) + + expect(result).toMatchObject({ + success: false, + data: { + notFound: [], + failed: [ + { id: KNOWLEDGE_BASE.id, name: KNOWLEDGE_BASE.name, reason: 'Knowledge base is locked' }, + ], }, - ]) - mockImportKnowledgeSearchResultSecretProvenance.mockResolvedValueOnce({ - imported: false, - documentMetadata: {}, }) + }) + + it('delegates document deletion and retains partial batch results', async () => { + mockDeleteKnowledgeDocument.mockRejectedValueOnce( + new OrchestrationError('not_found', 'Document not found') + ) const result = await knowledgeBaseServerTool.execute( { - operation: 'query', - args: { - knowledgeBaseId: 'knowledge-base-1', - query: 'public query', - }, + operation: 'delete_document', + args: { knowledgeBaseId: KNOWLEDGE_BASE.id, documentIds: ['missing', 'document-1'] }, }, - { - userId: 'external-admin', - workspaceId: 'workspace-paid', - toolCallId: 'tool-1', - copilotToolExecution: true, - billingAttribution: BILLING_ATTRIBUTION, - resolvedSecretTraceRegistry: registry, - } + CONTEXT ) - expect(result).toEqual({ - success: false, - message: 'Failed to query knowledge base: Knowledge result secret provenance is unavailable', + expect(result).toMatchObject({ + success: true, + data: { deleted: ['document-1'], failed: ['missing'] }, }) - expect(registry.isPermanentlyIncomplete()).toBe(true) + expectDelegatedPrincipal(mockDeleteKnowledgeDocument.mock.calls[1][0]) + expect(mockCaptureServerEvent).toHaveBeenCalledWith( + 'external-admin', + 'knowledge_base_document_deleted', + expect.objectContaining({ knowledge_base_id: KNOWLEDGE_BASE.id }), + expect.any(Object) + ) }) + + it.each([ + { + operation: 'add_file', + args: { knowledgeBaseId: KNOWLEDGE_BASE.id, filePaths: Array(101).fill('files/doc.pdf') }, + }, + { + operation: 'delete', + args: { knowledgeBaseIds: Array.from({ length: 101 }, (_, index) => `kb-${index}`) }, + }, + { + operation: 'delete_document', + args: { + knowledgeBaseId: KNOWLEDGE_BASE.id, + documentIds: Array.from({ length: 101 }, (_, index) => `document-${index}`), + }, + }, + ])( + 'rejects oversized $operation batches before application work', + async ({ operation, args }) => { + const result = await knowledgeBaseServerTool.execute({ operation, args }, CONTEXT) + + expect(result.success).toBe(false) + expect(result.message).toContain('Maximum is 100') + expect(mockReadKnowledgeBase).not.toHaveBeenCalled() + expect(mockDeleteKnowledgeBase).not.toHaveBeenCalled() + expect(mockDeleteKnowledgeDocument).not.toHaveBeenCalled() + expect(mockUploadKnowledgeDocument).not.toHaveBeenCalled() + } + ) }) -describe('knowledge base add_file usage gate', () => { +describe('knowledge_base add_file delegation', () => { beforeEach(() => { vi.clearAllMocks() - resetDbChainMock() - mockCheckKnowledgeBaseWriteAccess.mockResolvedValue({ - hasAccess: true, - knowledgeBase: { id: 'knowledge-base-1', workspaceId: 'workspace-paid', name: 'Paid KB' }, + mockReadKnowledgeBase.mockResolvedValue({ knowledgeBase: KNOWLEDGE_BASE, folderPath: '/' }) + mockResolveWorkspaceFileReference.mockResolvedValue({ + id: 'file-1', + key: 'workspace/workspace-paid/report.pdf', + name: 'report.pdf', + size: 100, + type: 'application/pdf', }) - vi.mocked(getKnowledgeBaseById).mockResolvedValue({ - id: 'knowledge-base-1', - workspaceId: 'workspace-paid', - } as Awaited>) - mockGetBoundWorkspaceFileSecretProvenance.mockResolvedValue({ - status: 'exact', - entries: [], + mockGetBoundWorkspaceFileSecretProvenance.mockResolvedValue({ status: 'exact', entries: [] }) + mockUploadKnowledgeDocument.mockResolvedValue({ + created: true, + document: { + id: 'document-1', + filename: 'report.pdf', + fileSize: 100, + mimeType: 'application/pdf', + }, }) }) - function addFile() { - return knowledgeBaseServerTool.execute( + it('preserves file resolution and performs current admission inside uploadKnowledgeDocument', async () => { + const result = await knowledgeBaseServerTool.execute( { operation: 'add_file', - args: { knowledgeBaseId: 'knowledge-base-1', filePaths: ['files/report.pdf'] }, + args: { knowledgeBaseId: KNOWLEDGE_BASE.id, filePaths: ['files/report.pdf'] }, }, - { - userId: 'external-admin', - workspaceId: 'workspace-paid', - toolCallId: 'tool-1', - copilotToolExecution: true, - billingAttribution: BILLING_ATTRIBUTION, - } + CONTEXT ) - } - - it('refuses to index when the payer is over its usage limit', async () => { - vi.mocked(checkAttributedUsageLimits).mockResolvedValue({ - isExceeded: true, - message: 'Usage limit exceeded.', - } as Awaited>) - - const result = await addFile() - - expect(result.success).toBe(false) - expect(result.message).toContain('Usage limit exceeded') - // The gate must precede any indexing work, matching the upload routes. - expect(resolveWorkspaceFileReference).not.toHaveBeenCalled() - expect(createSingleDocument).not.toHaveBeenCalled() - }) - it('gates on the knowledge base workspace payer, not the caller', async () => { - vi.mocked(checkAttributedUsageLimits).mockResolvedValue({ - isExceeded: false, - } as Awaited>) - vi.mocked(resolveWorkspaceFileReference).mockResolvedValue(null) - - await addFile() - - expect(checkAttributedUsageLimits).toHaveBeenCalledWith(BILLING_ATTRIBUTION) + expect(result).toMatchObject({ + success: true, + data: { added: [{ documentId: 'document-1', filename: 'report.pdf' }] }, + }) + expect(mockResolveWorkspaceFileReference).toHaveBeenCalledWith( + expect.objectContaining({ workspaceId: 'workspace-paid', reference: 'files/report.pdf' }) + ) + const call = mockUploadKnowledgeDocument.mock.calls[0][0] + expectDelegatedPrincipal(call) + expect(call.input).toMatchObject({ + knowledgeBaseId: KNOWLEDGE_BASE.id, + assertedWorkspaceId: 'workspace-paid', + startProcessing: true, + source: 'agent', + document: { filename: 'report.pdf', fileSize: 100, mimeType: 'application/pdf' }, + }) + expect(call.input).not.toHaveProperty('usageAdmission') + expect(mockKnowledgeBaseDocumentsUploaded).toHaveBeenCalledWith( + expect.objectContaining({ knowledgeBaseId: KNOWLEDGE_BASE.id, documentsCount: 1 }) + ) }) - it('does not index a workspace file containing resolved-secret provenance', async () => { - vi.mocked(checkAttributedUsageLimits).mockResolvedValue({ - isExceeded: false, - } as Awaited>) - vi.mocked(resolveWorkspaceFileReference).mockResolvedValue({ - id: 'file-1', - key: 'workspace/workspace-paid/report.pdf', - name: 'report.pdf', - size: 100, - type: 'application/pdf', - } as Awaited>) + it('rejects files carrying resolved-secret provenance before durable registration', async () => { mockGetBoundWorkspaceFileSecretProvenance.mockResolvedValueOnce({ status: 'exact', entries: [{ name: 'API_KEY', encryptedValue: 'encrypted-secret' }], }) - const result = await addFile() + const result = await knowledgeBaseServerTool.execute( + { + operation: 'add_file', + args: { knowledgeBaseId: KNOWLEDGE_BASE.id, filePaths: ['files/report.pdf'] }, + }, + CONTEXT + ) expect(result.success).toBe(false) - expect(mockGetBoundWorkspaceFileSecretProvenance).toHaveBeenCalledWith('workspace-paid', { - fileId: 'file-1', - key: 'workspace/workspace-paid/report.pdf', - context: 'workspace', - }) - expect(createSingleDocument).not.toHaveBeenCalled() + expect(mockUploadKnowledgeDocument).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.ts b/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.ts index c79a4c18da7..924e77186d0 100644 --- a/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.ts +++ b/apps/sim/lib/copilot/tools/server/knowledge/knowledge-base.ts @@ -6,12 +6,14 @@ import { generateId } from '@sim/utils/id' import { filterUndefined } from '@sim/utils/object' import { truncate } from '@sim/utils/string' import { and, eq, isNull } from 'drizzle-orm' -import { checkActorUsageLimits } from '@/lib/billing/calculations/usage-monitor' import { assertBillingAttributionSnapshot, type BillingAttributionSnapshot, - checkAttributedUsageLimits, } from '@/lib/billing/core/billing-attribution' +import { + messageForCopilotKnowledgeError, + resolveCopilotKnowledgePrincipal, +} from '@/lib/copilot/application/execute-knowledge-use-case' import { resolveCopilotFilePrincipal } from '@/lib/copilot/auth/file-delegation' import { KnowledgeBase } from '@/lib/copilot/generated/tool-catalog-v1' import { projectToolErrorMessageForCopilot } from '@/lib/copilot/request/tools/resolved-secret-result' @@ -22,25 +24,30 @@ import { } from '@/lib/copilot/tools/server/base-tool' import { projectServerToolModelInput } from '@/lib/copilot/tools/server/model-input' import { + asOrchestrationError, messageForOrchestrationError, type OrchestrationErrorCode, } from '@/lib/core/orchestration/types' -import { generateSearchEmbedding, recordSearchEmbeddingUsage } from '@/lib/knowledge/embeddings' +import { PlatformEvents } from '@/lib/core/telemetry' +import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/billing' +import { + deleteKnowledgeDocument, + uploadKnowledgeDocument, +} from '@/lib/knowledge/application/documents' +import { + createKnowledgeBase, + deleteKnowledgeBaseOperation, + readKnowledgeBase, + updateKnowledgeBaseOperation, +} from '@/lib/knowledge/application/knowledge-bases' +import { searchKnowledge } from '@/lib/knowledge/application/search' import { - performCreateKnowledgeBase, performCreateKnowledgeConnector, - performDeleteKnowledgeBase, performDeleteKnowledgeConnector, - performDeleteKnowledgeDocument, performSyncKnowledgeConnector, - performUpdateKnowledgeBase, performUpdateKnowledgeConnector, performUpdateKnowledgeDocument, - performUploadKnowledgeDocument, } from '@/lib/knowledge/orchestration' -import { executeKnowledgeSearch } from '@/lib/knowledge/search/queries' -import { importKnowledgeSearchResultSecretProvenance } from '@/lib/knowledge/secret-provenance' -import { getKnowledgeBaseById } from '@/lib/knowledge/service' import { createTagDefinition, deleteTagDefinition, @@ -50,6 +57,7 @@ import { getTagUsageStats, updateTagDefinition, } from '@/lib/knowledge/tags/service' +import { captureServerEvent } from '@/lib/posthog/server' import { StorageService } from '@/lib/uploads' import { getBoundWorkspaceFileSecretProvenance } from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance' import { fileOperations } from '@/lib/workspace-files/application/operations' @@ -62,6 +70,7 @@ import { } from '@/app/api/knowledge/utils' const logger = createLogger('KnowledgeBaseServerTool') +const MAX_COPILOT_KNOWLEDGE_BATCH_SIZE = 100 function requireKnowledgeBillingAttribution( context: ServerToolContext, @@ -91,6 +100,73 @@ function agentFacingError( return messageForOrchestrationError(outcome, fallback) } +function captureKnowledgeBaseCreated( + userId: string, + workspaceId: string, + knowledgeBase: { id: string; name: string } +): void { + PlatformEvents.knowledgeBaseCreated({ + knowledgeBaseId: knowledgeBase.id, + name: knowledgeBase.name, + workspaceId, + }) + captureServerEvent( + userId, + 'knowledge_base_created', + { + knowledge_base_id: knowledgeBase.id, + workspace_id: workspaceId, + name: knowledgeBase.name, + }, + { + groups: { workspace: workspaceId }, + setOnce: { first_kb_created_at: new Date().toISOString() }, + } + ) +} + +function captureKnowledgeDocumentUploaded( + userId: string, + workspaceId: string, + knowledgeBaseId: string, + document: { mimeType: string; fileSize: number } +): void { + PlatformEvents.knowledgeBaseDocumentsUploaded({ + knowledgeBaseId, + documentsCount: 1, + uploadType: 'single', + mimeType: document.mimeType, + fileSize: document.fileSize, + }) + captureServerEvent( + userId, + 'knowledge_base_document_uploaded', + { + knowledge_base_id: knowledgeBaseId, + workspace_id: workspaceId, + document_count: 1, + upload_type: 'single', + }, + { + groups: { workspace: workspaceId }, + setOnce: { first_document_uploaded_at: new Date().toISOString() }, + } + ) +} + +function captureKnowledgeDocumentDeleted( + userId: string, + workspaceId: string, + knowledgeBaseId: string +): void { + captureServerEvent( + userId, + 'knowledge_base_document_deleted', + { knowledge_base_id: knowledgeBaseId, workspace_id: workspaceId }, + { groups: { workspace: workspaceId } } + ) +} + type KnowledgeBaseArgs = { operation: string args?: Record @@ -111,17 +187,10 @@ export const knowledgeBaseServerTool: BaseServerTool { - const withMessageId = (message: string) => - context?.messageId ? `${message} [messageId:${context.messageId}]` : message - - if (!context?.userId) { - logger.error('Unauthorized attempt to access knowledge base - no authenticated user context') - throw new Error('Authentication required') - } - + if (!context) throw new Error('Knowledge delegation requires a Copilot execution context') + const principal = resolveCopilotKnowledgePrincipal(context) const { operation, args = {} } = params - const workspaceId = - context.workspaceId || ((args as Record).workspaceId as string | undefined) + const workspaceId = principal.workspaceId const assertNotAborted = () => assertServerToolNotAborted( context, @@ -156,23 +225,18 @@ export const knowledgeBaseServerTool: BaseServerTool MAX_COPILOT_KNOWLEDGE_BATCH_SIZE) { return { success: false, - message: `Knowledge base with ID "${args.knowledgeBaseId}" not found`, + message: `Too many files (${fileRefs.length}). Maximum is ${MAX_COPILOT_KNOWLEDGE_BATCH_SIZE}.`, } } - const targetKb = await getKnowledgeBaseById(args.knowledgeBaseId) - if (!targetKb || !targetKb.workspaceId) { - return { - success: false, - message: `Knowledge base with ID "${args.knowledgeBaseId}" not found`, - } - } - - const kbWorkspaceId: string = targetKb.workspaceId - const billingAttribution = requireKnowledgeBillingAttribution(context, kbWorkspaceId) - - // Gate the payer before accepting indexing work, same as the upload routes. - const usage = await checkAttributedUsageLimits(billingAttribution) - if (usage.isExceeded) { - return { - success: false, - message: - usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.', - } - } + const { knowledgeBase: targetKb } = await readKnowledgeBase.execute({ + principal, + input: { + knowledgeBaseId: args.knowledgeBaseId, + assertedWorkspaceId: workspaceId, + }, + }) const added: Array<{ documentId: string; filename: string }> = [] const failedFiles: string[] = [] @@ -405,15 +403,19 @@ export const knowledgeBaseServerTool: BaseServerTool a.filename).join(', ') @@ -496,33 +509,16 @@ export const knowledgeBaseServerTool: BaseServerTool MAX_COPILOT_KNOWLEDGE_BATCH_SIZE) { + return { + success: false, + message: `Too many knowledge base IDs (${kbIds.length}). Maximum is ${MAX_COPILOT_KNOWLEDGE_BATCH_SIZE}.`, + } + } const deleted: Array<{ id: string; name: string }> = [] const notFound: string[] = [] @@ -555,41 +557,42 @@ export const knowledgeBaseServerTool: BaseServerTool = [] for (const kbId of kbIds) { - const writeAccess = await checkKnowledgeBaseWriteAccess(kbId, context.userId) - if (!writeAccess.hasAccess) { - notFound.push(kbId) - continue - } - - const kbToDelete = await getKnowledgeBaseById(kbId) - if (!kbToDelete) { - notFound.push(kbId) - continue - } - - const requestId = generateId().slice(0, 8) - assertNotAborted() - const outcome = await performDeleteKnowledgeBase({ - ...actor(requestId), - knowledgeBase: { - id: kbId, - name: kbToDelete.name, - workspaceId: kbToDelete.workspaceId, - }, - }) - if (!outcome.success) { - if (outcome.errorCode === 'not_found') { + let knowledgeBaseName = kbId + try { + const readResult = await readKnowledgeBase.execute({ + principal, + input: { knowledgeBaseId: kbId, assertedWorkspaceId: workspaceId }, + }) + knowledgeBaseName = readResult.knowledgeBase.name + assertNotAborted() + const deletedKnowledgeBase = await deleteKnowledgeBaseOperation.execute({ + principal, + input: { + knowledgeBaseId: kbId, + assertedWorkspaceId: workspaceId, + source: 'agent', + }, + }) + PlatformEvents.knowledgeBaseDeleted({ knowledgeBaseId: kbId }) + deleted.push(deletedKnowledgeBase) + } catch (error) { + const classified = asOrchestrationError(error) + if ( + classified?.code === 'not_found' || + classified?.code === 'forbidden' || + classified?.code === 'unauthorized' + ) { notFound.push(kbId) - } else { + } else if (classified && classified.code !== 'internal') { failed.push({ id: kbId, - name: kbToDelete.name, - reason: agentFacingError(outcome, 'Failed to delete knowledge base'), + name: knowledgeBaseName, + reason: classified.message, }) + } else { + throw error } - continue } - deleted.push({ id: kbId, name: kbToDelete.name }) } const deleteSummary = [ @@ -619,35 +622,37 @@ export const knowledgeBaseServerTool: BaseServerTool MAX_COPILOT_KNOWLEDGE_BATCH_SIZE) { + return { + success: false, + message: `Too many document IDs (${docIds.length}). Maximum is ${MAX_COPILOT_KNOWLEDGE_BATCH_SIZE}.`, + } + } const deleted: string[] = [] const failed: string[] = [] for (const docId of docIds) { assertNotAborted() - const docAccess = await checkDocumentWriteAccess( - args.knowledgeBaseId, - docId, - context.userId - ) - if (!docAccess.hasAccess) { - failed.push(docId) - continue - } - const requestId = generateId().slice(0, 8) - const outcome = await performDeleteKnowledgeDocument({ - ...actor(requestId), - knowledgeBase: { - id: args.knowledgeBaseId, - name: docAccess.knowledgeBase.name, - workspaceId: docAccess.knowledgeBase.workspaceId ?? null, - }, - document: docAccess.document, - }) - if (outcome.success) { + try { + await deleteKnowledgeDocument.execute({ + principal, + input: { + knowledgeBaseId: args.knowledgeBaseId, + documentId: docId, + assertedWorkspaceId: workspaceId, + source: 'agent', + }, + }) + captureKnowledgeDocumentDeleted(context.userId, workspaceId, args.knowledgeBaseId) deleted.push(docId) - } else { - failed.push(docId) + } catch (error) { + const classified = asOrchestrationError(error) + if (classified && classified.code !== 'internal') { + failed.push(docId) + continue + } + throw error } } @@ -1206,9 +1211,33 @@ export const knowledgeBaseServerTool: BaseServerTool ({ executeKnowledgeBase: vi.fn() })) + +vi.mock('@/lib/copilot/generated/tool-catalog-v1', () => ({ + SearchKnowledgeBase: { id: 'search_knowledge_base' }, +})) +vi.mock('@/lib/copilot/tools/server/knowledge/knowledge-base', () => ({ + knowledgeBaseServerTool: { execute: executeKnowledgeBase }, +})) + +import { searchKnowledgeBaseServerTool } from '@/lib/copilot/tools/server/knowledge/search-knowledge-base' + +describe('search_knowledge_base delegation', () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it.each(['get', 'query'])('forwards %s with the immutable trusted context', async (operation) => { + const params = { operation, args: { knowledgeBaseId: 'kb-1' } } + const context = { + userId: 'user-1', + workspaceId: 'workspace-1', + chatId: 'chat-1', + toolCallId: 'tool-1', + copilotToolExecution: true, + } + executeKnowledgeBase.mockResolvedValueOnce({ success: true, message: 'ok' }) + + await expect(searchKnowledgeBaseServerTool.execute(params, context)).resolves.toEqual({ + success: true, + message: 'ok', + }) + expect(executeKnowledgeBase).toHaveBeenCalledWith(params, context) + }) + + it('does not expose the legacy delete compatibility operation', async () => { + const result = await searchKnowledgeBaseServerTool.execute( + { operation: 'delete', args: { knowledgeBaseId: 'kb-1' } }, + { + userId: 'user-1', + workspaceId: 'workspace-1', + toolCallId: 'tool-1', + copilotToolExecution: true, + } + ) + + expect(result.success).toBe(false) + expect(result.message).toContain('read-only') + expect(executeKnowledgeBase).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/copilot/vfs/workspace-vfs.ts b/apps/sim/lib/copilot/vfs/workspace-vfs.ts index b451518ff56..1458b9249d4 100644 --- a/apps/sim/lib/copilot/vfs/workspace-vfs.ts +++ b/apps/sim/lib/copilot/vfs/workspace-vfs.ts @@ -4,7 +4,6 @@ import { db } from '@sim/db' import { chat as chatTable, customTools as customToolsTable, - document, folder as folderTable, knowledgeBaseTagDefinitions, knowledgeConnector, @@ -118,7 +117,9 @@ import { isOAuthServiceDeploymentAvailable, } from '@/lib/integrations/availability.server' import { createIntegrationCredentialVisibility } from '@/lib/integrations/credential-visibility.server' -import { getKnowledgeBases } from '@/lib/knowledge/service' +import { listKnowledgeDocuments } from '@/lib/knowledge/application/documents' +import { listKnowledgeBases } from '@/lib/knowledge/application/knowledge-bases' +import { getKnowledgeBases as getLegacyKnowledgeBases } from '@/lib/knowledge/service' import { validateMermaidSource } from '@/lib/mermaid/validate' import { isBlockTypeAccessControlExempt } from '@/lib/permission-groups/block-access' import { intersectIntegrationAllowlists } from '@/lib/permission-groups/integration-allowlist' @@ -157,6 +158,8 @@ const logger = createLogger('WorkspaceVFS') // double-cast-allowed: a no-op stands in for the unused SVG-typed BlockIcon slot const PLACEHOLDER_BLOCK_ICON = (() => null) as unknown as BlockIcon const MAX_COMPILED_ATTACHMENT_BYTES = 5 * 1024 * 1024 +const KNOWLEDGE_DOCUMENT_PAGE_SIZE = 100 +const MAX_VFS_KNOWLEDGE_DOCUMENTS = 10_000 function bindWorkspaceFileResult( record: WorkspaceFileRecord, @@ -566,6 +569,7 @@ function getStaticComponentFiles(): Map { */ export class WorkspaceVFS { private readonly filePrincipal?: Principal + private readonly knowledgePrincipal?: Principal // Eagerly-materialized, cheap content (structure + metadata): folder markers, // per-resource meta.json, WORKSPACE.md/WORKSPACE_CONTEXT.md, static components. private files: Map = new Map() @@ -598,8 +602,9 @@ export class WorkspaceVFS { */ private _customBlockTypes: Set | null = null - constructor(filePrincipal?: Principal) { + constructor(filePrincipal?: Principal, knowledgePrincipal?: Principal) { this.filePrincipal = filePrincipal + this.knowledgePrincipal = knowledgePrincipal } get workspaceId(): string { @@ -799,7 +804,7 @@ export class WorkspaceVFS { sandboxEntitled, ] = await Promise.all([ timed('workflows', this.materializeWorkflows(workspaceId)), - timed('knowledge_bases', this.materializeKnowledgeBases(workspaceId, userId)), + timed('knowledge_bases', this.materializeKnowledgeBases(workspaceId)), timed('tables', this.materializeTables(workspaceId)), timed('files', this.materializeFiles(workspaceId)), timed( @@ -1738,99 +1743,120 @@ export class WorkspaceVFS { })) } - /** - * Materialize knowledge bases using the shared getKnowledgeBases function. - * Returns a summary for WORKSPACE.md generation. - */ + /** Materializes authorized knowledge summaries for WORKSPACE.md generation. */ private async materializeKnowledgeBases( - workspaceId: string, - userId: string + workspaceId: string ): Promise { - const kbs = await getKnowledgeBases(userId, workspaceId) + if (!this.knowledgePrincipal) { + throw new Error('Workspace VFS knowledge materialization requires a trusted principal') + } + const { knowledgeBases } = await listKnowledgeBases.execute({ + principal: this.knowledgePrincipal, + input: { workspaceId }, + }) + const kbs = knowledgeBases.map(({ knowledgeBase }) => knowledgeBase) const tagDefinitionsByKb = await this.loadKbTagDefinitions(kbs.map((kb) => kb.id)) - await Promise.all( - kbs.map(async (kb) => { - const safeName = sanitizeName(kb.name) - const prefix = `knowledgebases/${safeName}/` + for (const kb of kbs) { + const safeName = sanitizeName(kb.name) + const prefix = `knowledgebases/${safeName}/` - this.files.set( - `${prefix}meta.json`, - serializeKBMeta({ - id: kb.id, - name: kb.name, - description: kb.description, - embeddingModel: kb.embeddingModel, - embeddingDimension: kb.embeddingDimension, - tokenCount: kb.tokenCount, - createdAt: kb.createdAt, - updatedAt: kb.updatedAt, - documentCount: kb.docCount, - connectorTypes: kb.connectorTypes, - tagDefinitions: tagDefinitionsByKb.get(kb.id), - }) - ) + this.files.set( + `${prefix}meta.json`, + serializeKBMeta({ + id: kb.id, + name: kb.name, + description: kb.description, + embeddingModel: kb.embeddingModel, + embeddingDimension: kb.embeddingDimension, + tokenCount: kb.tokenCount, + createdAt: kb.createdAt, + updatedAt: kb.updatedAt, + documentCount: kb.docCount, + connectorTypes: kb.connectorTypes, + tagDefinitions: tagDefinitionsByKb.get(kb.id), + }) + ) - // documents.json / connectors.json are lazy, advertised only when the KB - // summary says they exist (docCount / connectorTypes) — no per-KB query on - // a read/glob, only when the artifact is read or grepped. - if (kb.docCount > 0) { - this.registerLazy(`${prefix}documents.json`, async () => { - const docRows = await db - .select({ - id: document.id, - filename: document.filename, - fileSize: document.fileSize, - mimeType: document.mimeType, - chunkCount: document.chunkCount, - tokenCount: document.tokenCount, - processingStatus: document.processingStatus, - enabled: document.enabled, - uploadedAt: document.uploadedAt, - }) - .from(document) - .where( - and( - eq(document.knowledgeBaseId, kb.id), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt) - ) + // documents.json / connectors.json are lazy, advertised only when the KB + // summary says they exist (docCount / connectorTypes) — no per-KB query on + // a read/glob, only when the artifact is read or grepped. + if (kb.docCount > 0) { + this.registerLazy(`${prefix}documents.json`, async () => { + if (!this.knowledgePrincipal) { + throw new Error('Workspace VFS knowledge document read requires a trusted principal') + } + if (kb.docCount > MAX_VFS_KNOWLEDGE_DOCUMENTS) { + throw new Error( + `Knowledge base ${kb.id} has more than ${MAX_VFS_KNOWLEDGE_DOCUMENTS} documents; documents.json cannot be materialized` + ) + } + const documents: Awaited>['documents'] = + [] + let offset = 0 + while (true) { + const page = await listKnowledgeDocuments.execute({ + principal: this.knowledgePrincipal, + input: { + knowledgeBaseId: kb.id, + assertedWorkspaceId: workspaceId, + limit: KNOWLEDGE_DOCUMENT_PAGE_SIZE, + offset, + }, + }) + documents.push(...page.documents) + if (documents.length > MAX_VFS_KNOWLEDGE_DOCUMENTS) { + throw new Error( + `Knowledge base ${kb.id} exceeded the ${MAX_VFS_KNOWLEDGE_DOCUMENTS} document limit while materializing documents.json` ) - return docRows.length > 0 ? serializeDocuments(docRows) : null - }) - } + } + if (!page.pagination.hasMore) break + offset += page.pagination.limit + } + const docRows = documents.map((document) => ({ + id: document.id, + filename: document.filename, + fileSize: document.fileSize, + mimeType: document.mimeType, + chunkCount: document.chunkCount, + tokenCount: document.tokenCount, + processingStatus: document.processingStatus, + enabled: document.enabled, + uploadedAt: document.uploadedAt, + })) + return docRows.length > 0 ? serializeDocuments(docRows) : null + }) + } - if (kb.connectorTypes.length > 0) { - this.registerLazy(`${prefix}connectors.json`, async () => { - const connectorRows = await db - .select({ - id: knowledgeConnector.id, - connectorType: knowledgeConnector.connectorType, - status: knowledgeConnector.status, - syncMode: knowledgeConnector.syncMode, - syncIntervalMinutes: knowledgeConnector.syncIntervalMinutes, - lastSyncAt: knowledgeConnector.lastSyncAt, - lastSyncError: knowledgeConnector.lastSyncError, - lastSyncDocCount: knowledgeConnector.lastSyncDocCount, - nextSyncAt: knowledgeConnector.nextSyncAt, - consecutiveFailures: knowledgeConnector.consecutiveFailures, - createdAt: knowledgeConnector.createdAt, - }) - .from(knowledgeConnector) - .where( - and( - eq(knowledgeConnector.knowledgeBaseId, kb.id), - isNull(knowledgeConnector.archivedAt), - isNull(knowledgeConnector.deletedAt) - ) + if (kb.connectorTypes.length > 0) { + this.registerLazy(`${prefix}connectors.json`, async () => { + const connectorRows = await db + .select({ + id: knowledgeConnector.id, + connectorType: knowledgeConnector.connectorType, + status: knowledgeConnector.status, + syncMode: knowledgeConnector.syncMode, + syncIntervalMinutes: knowledgeConnector.syncIntervalMinutes, + lastSyncAt: knowledgeConnector.lastSyncAt, + lastSyncError: knowledgeConnector.lastSyncError, + lastSyncDocCount: knowledgeConnector.lastSyncDocCount, + nextSyncAt: knowledgeConnector.nextSyncAt, + consecutiveFailures: knowledgeConnector.consecutiveFailures, + createdAt: knowledgeConnector.createdAt, + }) + .from(knowledgeConnector) + .where( + and( + eq(knowledgeConnector.knowledgeBaseId, kb.id), + isNull(knowledgeConnector.archivedAt), + isNull(knowledgeConnector.deletedAt) ) - return connectorRows.length > 0 ? serializeConnectors(connectorRows) : null - }) - } - }) - ) + ) + return connectorRows.length > 0 ? serializeConnectors(connectorRows) : null + }) + } + } return kbs.map((kb) => ({ id: kb.id, @@ -2343,7 +2369,7 @@ export class WorkspaceVFS { input: { workspaceId, scope: 'archived' }, }) .then(({ folders }) => folders), - getKnowledgeBases(userId, workspaceId, 'archived'), + getLegacyKnowledgeBases(userId, workspaceId, 'archived'), ]) for (const wf of archivedWorkflows) { @@ -2570,10 +2596,14 @@ export class WorkspaceVFS { export async function getOrMaterializeVFS( workspaceId: string, userId: string, - options?: { secretMountPolicy?: SecretMountPolicy; filePrincipal?: Principal } + options?: { + secretMountPolicy?: SecretMountPolicy + filePrincipal?: Principal + knowledgePrincipal?: Principal + } ): Promise { await assertActiveWorkspaceAccess(workspaceId, userId) - const vfs = new WorkspaceVFS(options?.filePrincipal) + const vfs = new WorkspaceVFS(options?.filePrincipal, options?.knowledgePrincipal) await vfs.materialize(workspaceId, userId, options) return vfs } diff --git a/apps/sim/lib/folders/cascade.test.ts b/apps/sim/lib/folders/cascade.test.ts index f354a9ded09..e75fa3ee45d 100644 --- a/apps/sim/lib/folders/cascade.test.ts +++ b/apps/sim/lib/folders/cascade.test.ts @@ -41,7 +41,11 @@ function makeTx(options: { selects?: unknown[][]; updates?: unknown[][] } = {}) from: () => ({ where: (where: unknown) => { selectCalls.push({ where }) - return Promise.resolve(selectQueue.shift() ?? []) + const rows = selectQueue.shift() ?? [] + return { + limit: () => Promise.resolve(rows), + then: (resolve: (value: unknown) => unknown) => Promise.resolve(rows).then(resolve), + } }, }), }), @@ -139,6 +143,22 @@ describe('collectCascadeSubtreeIds', () => { ) expect(hasMockCondition(selectCalls[0].where, (node) => node.right === 'ws-1')).toBe(true) }) + + it('fails before materializing an oversized recursive cascade', async () => { + const { tx } = makeTx({ + selects: [ + [ + { id: 'root', parentId: null }, + { id: 'child', parentId: 'root' }, + { id: 'grandchild', parentId: 'child' }, + ], + ], + }) + + await expect( + collectCascadeSubtreeIds(tx, 'ws-1', 'knowledge_base', 'root', TIMESTAMP, 2) + ).rejects.toThrow('Folder cascade exceeds the 2 row limit') + }) }) describe('collectArchivedSubtreeIds', () => { diff --git a/apps/sim/lib/folders/cascade.ts b/apps/sim/lib/folders/cascade.ts index 3e72aa4fd02..c4da00dedd7 100644 --- a/apps/sim/lib/folders/cascade.ts +++ b/apps/sim/lib/folders/cascade.ts @@ -32,9 +32,10 @@ export async function collectCascadeSubtreeIds( workspaceId: string, resourceType: FolderResourceType, folderId: string, - timestamp: Date + timestamp: Date, + maxRows?: number ): Promise { - const cascadeFolders = await tx + const query = tx .select({ id: folderTable.id, parentId: folderTable.parentId }) .from(folderTable) .where( @@ -44,6 +45,10 @@ export async function collectCascadeSubtreeIds( or(isNull(folderTable.deletedAt), eq(folderTable.deletedAt, timestamp)) ) ) + const cascadeFolders = maxRows === undefined ? await query : await query.limit(maxRows + 1) + if (maxRows !== undefined && cascadeFolders.length > maxRows) { + throw new Error(`Folder cascade exceeds the ${maxRows} row limit`) + } return [folderId, ...collectDescendantFolderIds(cascadeFolders, folderId)] } diff --git a/apps/sim/lib/folders/orchestration.ts b/apps/sim/lib/folders/orchestration.ts index 582b1f755cd..f2540f15bff 100644 --- a/apps/sim/lib/folders/orchestration.ts +++ b/apps/sim/lib/folders/orchestration.ts @@ -69,6 +69,7 @@ export interface DeleteFolderParams { userId: string folderName?: string folderPath?: string + maxFolderRows?: number } export interface DeleteFolderResult { @@ -103,6 +104,9 @@ export interface DeleteFolderByPathParams { userId: string path: string recursive: boolean + maxFolderRows?: number + effects?: boolean + throwInfrastructure?: boolean } export interface DeleteFolderByPathResult extends DeleteFolderResult { @@ -164,7 +168,12 @@ function pathMutationError(error: unknown): FolderPathMutationResult { } async function executeCreateFolderAtPath( - params: Omit & { path: string }, + params: Omit & { + path: string + effects?: boolean + throwInfrastructure?: boolean + maxFolderRows?: number + }, projectLegacyLifecycle: boolean ): Promise { try { @@ -173,7 +182,9 @@ async function executeCreateFolderAtPath( const folder = await withTransactionRetry( async (tx) => { await acquireFolderMutationLock(tx, params.workspaceId, params.resourceType) - const index = await loadActiveFolderPathIndex(params.workspaceId, params.resourceType, tx) + const index = await loadActiveFolderPathIndex(params.workspaceId, params.resourceType, tx, { + maxRows: params.maxFolderRows, + }) if (index.idByPath.has(params.path)) throw new Error(DUPLICATE_NAME_ERROR) const parentPath = parentFolderPath(params.path) @@ -213,7 +224,7 @@ async function executeCreateFolderAtPath( { label: 'create-folder-at-path' } ) - if (projectLegacyLifecycle) { + if (projectLegacyLifecycle && params.effects !== false) { recordAudit({ workspaceId: params.workspaceId, actorId: params.userId, @@ -225,21 +236,30 @@ async function executeCreateFolderAtPath( metadata: { path: params.path, folderResourceType: params.resourceType }, }) } - await notifyFolderResourceChanged(params.resourceType, params.workspaceId) + if (params.effects !== false) { + await notifyFolderResourceChanged(params.resourceType, params.workspaceId) + } return { success: true, folder, path: params.path } } catch (error) { - return pathMutationError(error) + const result = pathMutationError(error) + if (params.throwInfrastructure && result.errorCode === 'internal') throw error + return result } } /** Creates exactly the leaf identified by `path`; every ancestor must already exist. */ export async function createFolderAtPath( - params: Omit & { path: string } + params: Omit & { + path: string + effects?: boolean + throwInfrastructure?: boolean + maxFolderRows?: number + } ): Promise { return executeCreateFolderAtPath(params, true) } -/** Applies the authoritative mutation without projecting audit or realtime side effects. */ +/** Applies the authoritative mutation without projecting legacy audit. */ export async function createFolderAtPathTransition( params: Omit & { path: string } ): Promise { @@ -252,6 +272,9 @@ type RelocateFolderByPathParams = { userId: string path: string destinationPath: string + effects?: boolean + throwInfrastructure?: boolean + maxFolderRows?: number } async function executeRelocateFolderByPath( @@ -266,7 +289,9 @@ async function executeRelocateFolderByPath( const folder = await withTransactionRetry( async (tx) => { await acquireFolderMutationLock(tx, params.workspaceId, params.resourceType) - const index = await loadActiveFolderPathIndex(params.workspaceId, params.resourceType, tx) + const index = await loadActiveFolderPathIndex(params.workspaceId, params.resourceType, tx, { + maxRows: params.maxFolderRows, + }) const folderId = resolveRequiredFolderId(index, params.path) if (index.idByPath.has(params.destinationPath)) throw new Error(DUPLICATE_NAME_ERROR) @@ -309,7 +334,7 @@ async function executeRelocateFolderByPath( { label: 'relocate-folder-by-path' } ) - if (projectLegacyLifecycle) { + if (projectLegacyLifecycle && params.effects !== false) { recordAudit({ workspaceId: params.workspaceId, actorId: params.userId, @@ -325,10 +350,14 @@ async function executeRelocateFolderByPath( }, }) } - await notifyFolderResourceChanged(params.resourceType, params.workspaceId) + if (params.effects !== false) { + await notifyFolderResourceChanged(params.resourceType, params.workspaceId) + } return { success: true, folder, path: params.destinationPath } } catch (error) { - return pathMutationError(error) + const result = pathMutationError(error) + if (params.throwInfrastructure && result.errorCode === 'internal') throw error + return result } } @@ -339,7 +368,7 @@ export async function relocateFolderByPath( return executeRelocateFolderByPath(params, true) } -/** Applies the authoritative mutation without projecting audit or realtime side effects. */ +/** Applies the authoritative mutation without projecting legacy audit. */ export async function relocateFolderByPathTransition( params: RelocateFolderByPathParams ): Promise { @@ -356,7 +385,9 @@ async function executeDeleteFolderByPath( params.workspaceId, params.resourceType, async (tx) => { - const index = await loadActiveFolderPathIndex(params.workspaceId, params.resourceType, tx) + const index = await loadActiveFolderPathIndex(params.workspaceId, params.resourceType, tx, { + maxRows: params.maxFolderRows, + }) const folderId = resolveRequiredFolderId(index, params.path) if ( folderResourceConfig(params.resourceType).supportsLocking && @@ -394,11 +425,16 @@ async function executeDeleteFolderByPath( userId: params.userId, folderName: row.name, folderPath: params.path, + maxFolderRows: params.maxFolderRows, } } ) - const result = await deleteFolderWithoutTreeLock(resolved, null, projectLegacyLifecycle) + const effects = params.effects !== false + const result = await deleteFolderWithoutTreeLock(resolved, null, { + projectAudit: projectLegacyLifecycle && effects, + notify: effects, + }) return { ...result, path: result.success ? params.path : undefined, @@ -406,7 +442,9 @@ async function executeDeleteFolderByPath( folderName: result.success ? resolved.folderName : undefined, } } catch (error) { - return pathMutationError(error) + const result = pathMutationError(error) + if (params.throwInfrastructure && result.errorCode === 'internal') throw error + return result } } @@ -417,7 +455,7 @@ export async function deleteFolderByPath( return executeDeleteFolderByPath(params, true) } -/** Applies the authoritative mutation without projecting audit or realtime side effects. */ +/** Applies the authoritative mutation without projecting legacy audit. */ export async function deleteFolderByPathTransition( params: DeleteFolderByPathParams ): Promise { @@ -726,13 +764,16 @@ export async function deleteFolder(params: DeleteFolderParams): Promise { const { resourceType, folderId, workspaceId, userId, folderName, folderPath } = params const config = folderResourceConfig(resourceType) @@ -741,13 +782,17 @@ async function deleteFolderWithoutTreeLock( // it is what distinguishes folders this cascade already stamped from folders archived // independently. const timestamp = deletedAt ?? new Date() - const folderIds = await collectCascadeSubtreeIds( - db, - workspaceId, - resourceType, - folderId, - timestamp - ) + const folderIds = + params.maxFolderRows === undefined + ? await collectCascadeSubtreeIds(db, workspaceId, resourceType, folderId, timestamp) + : await collectCascadeSubtreeIds( + db, + workspaceId, + resourceType, + folderId, + timestamp, + params.maxFolderRows + ) const rejection = await config.guardDelete?.({ workspaceId, folderIds }) if (rejection) { @@ -758,7 +803,7 @@ async function deleteFolderWithoutTreeLock( logger.info('Deleted folder and all contents', { folderId, resourceType, counts }) - if (projectLegacyLifecycle) { + if (options.projectAudit) { recordAudit({ workspaceId, actorId: userId, @@ -778,7 +823,7 @@ async function deleteFolderWithoutTreeLock( }) } // Live resource list (e.g. tables): a delete removes the folder and cascades to its contents. - await notifyFolderResourceChanged(resourceType, workspaceId) + if (options.notify) await notifyFolderResourceChanged(resourceType, workspaceId) return { success: true, deletedItems: toCascadeCounts(config, counts) } } diff --git a/apps/sim/lib/folders/queries.test.ts b/apps/sim/lib/folders/queries.test.ts index d6210ed53b5..3d7a8e1ef02 100644 --- a/apps/sim/lib/folders/queries.test.ts +++ b/apps/sim/lib/folders/queries.test.ts @@ -11,7 +11,9 @@ import { import { beforeEach, describe, expect, it, vi } from 'vitest' import { findActiveFolder, + listActiveFolderRows, listFoldersForWorkspace, + loadActiveFolderPathIndex, resolveRestoredFolderId, toFolderApi, wouldCreateFolderCycle, @@ -184,6 +186,26 @@ describe('folder queries', () => { }) }) + describe('bounded folder reads', () => { + it('fails before building an oversized path index', async () => { + queueTableRows(schemaMock.folder, [ROW, { ...ROW, id: 'f-2' }, { ...ROW, id: 'f-3' }]) + + await expect( + loadActiveFolderPathIndex('ws-1', 'knowledge_base', undefined, { maxRows: 2 }) + ).rejects.toThrow('Folder path index exceeds the 2 row limit') + expect(dbChainMockFns.limit).toHaveBeenCalledWith(3) + }) + + it('fails before returning an oversized folder list', async () => { + queueTableRows(schemaMock.folder, [ROW, { ...ROW, id: 'f-2' }, { ...ROW, id: 'f-3' }]) + + await expect(listActiveFolderRows('ws-1', 'knowledge_base', { maxRows: 2 })).rejects.toThrow( + 'Folder list exceeds the 2 row limit' + ) + expect(dbChainMockFns.limit).toHaveBeenCalledWith(3) + }) + }) + describe('toFolderApi', () => { it('serializes timestamps to ISO strings and preserves a null deletedAt', () => { expect(toFolderApi(ROW)).toMatchObject({ diff --git a/apps/sim/lib/folders/queries.ts b/apps/sim/lib/folders/queries.ts index 7811476e82d..ff2bade9564 100644 --- a/apps/sim/lib/folders/queries.ts +++ b/apps/sim/lib/folders/queries.ts @@ -166,14 +166,16 @@ interface ListActiveFolderRowsOptions { search?: string sortBy?: Exclude sortOrder?: V2SortOrder + maxRows?: number } export async function loadActiveFolderPathIndex( workspaceId: string, resourceType: FolderResourceType, - tx: DbOrTx = db + tx: DbOrTx = db, + options?: { maxRows?: number } ): Promise> { - const rows = await tx + const query = tx .select() .from(folder) .where( @@ -183,6 +185,10 @@ export async function loadActiveFolderPathIndex( isNull(folder.deletedAt) ) ) + const rows = options?.maxRows === undefined ? await query : await query.limit(options.maxRows + 1) + if (options?.maxRows !== undefined && rows.length > options.maxRows) { + throw new Error(`Folder path index exceeds the ${options.maxRows} row limit`) + } return buildFolderPathIndex(rows) } @@ -208,7 +214,7 @@ export async function listActiveFolderRows( ? isNull(folder.parentId) : eq(folder.parentId, options.parentId) - return tx + const query = tx .select() .from(folder) .where( @@ -221,6 +227,11 @@ export async function listActiveFolderRows( ) ) .orderBy(...listOrderBy(FOLDER_SORTS[options.sortBy ?? 'name'], options.sortOrder ?? 'asc')) + const rows = options.maxRows === undefined ? await query : await query.limit(options.maxRows + 1) + if (options.maxRows !== undefined && rows.length > options.maxRows) { + throw new Error(`Folder list exceeds the ${options.maxRows} row limit`) + } + return rows } /** diff --git a/apps/sim/lib/knowledge/application/authorization.test.ts b/apps/sim/lib/knowledge/application/authorization.test.ts new file mode 100644 index 00000000000..4952e6f787b --- /dev/null +++ b/apps/sim/lib/knowledge/application/authorization.test.ts @@ -0,0 +1,55 @@ +/** + * @vitest-environment node + */ + +import type { DelegatedPrincipal } from '@sim/auth/principal' +import { describe, expect, it } from 'vitest' +import { + KNOWLEDGE_DELEGATION_AUDIENCE, + knowledgeDelegationPolicy, +} from '@/lib/knowledge/application/authorization' +import { createKnowledgeDelegatedPrincipal } from '@/lib/knowledge/application/delegated-principal' + +describe('knowledge delegation policy', () => { + it('binds trusted delegation to the canonical workspace and audience', () => { + const principal = createKnowledgeDelegatedPrincipal({ + serviceId: 'copilot', + subjectUserId: 'user-1', + workspaceId: 'workspace-1', + delegationId: 'tool-call-1', + chatId: 'chat-1', + }) + + expect(principal.audience).toBe(KNOWLEDGE_DELEGATION_AUDIENCE) + expect(principal.resourceScope).toEqual({ chatId: 'chat-1' }) + expect( + knowledgeDelegationPolicy.isWithinScope(principal, { + workspaceId: 'workspace-1', + workspaceOrganizationId: null, + allowPersonalApiKeys: true, + }) + ).toBe(true) + expect( + knowledgeDelegationPolicy.isWithinScope(principal, { + workspaceId: 'workspace-2', + workspaceOrganizationId: null, + allowPersonalApiKeys: true, + }) + ).toBe(false) + }) + + it('does not accept a model-authored audience', () => { + const principal: DelegatedPrincipal = { + kind: 'delegated', + serviceId: 'copilot', + subjectUserId: 'user-1', + workspaceId: 'workspace-1', + delegationId: 'tool-call-1', + audience: 'model:chosen', + issuedAt: new Date(), + expiresAt: new Date(Date.now() + 60_000), + } + + expect(principal.audience).not.toBe(knowledgeDelegationPolicy.audience) + }) +}) diff --git a/apps/sim/lib/knowledge/application/authorization.ts b/apps/sim/lib/knowledge/application/authorization.ts new file mode 100644 index 00000000000..e75d1c496d4 --- /dev/null +++ b/apps/sim/lib/knowledge/application/authorization.ts @@ -0,0 +1,27 @@ +import type { Principal } from '@sim/auth/principal' +import type { + WorkspaceAuthorizationContext, + WorkspaceAuthorizationOptions, +} from '@/lib/core/application' + +export const KNOWLEDGE_DELEGATION_AUDIENCE = 'sim:knowledge' + +export interface KnowledgeAuthorizationContext extends WorkspaceAuthorizationContext { + knowledgeBaseId?: string + documentId?: string +} + +export type KnowledgeAuthorizationOptions = Omit< + WorkspaceAuthorizationOptions, + 'delegation' +> + +export const knowledgeDelegationPolicy = { + audience: KNOWLEDGE_DELEGATION_AUDIENCE, + isWithinScope( + delegated: Extract, + canonicalContext: KnowledgeAuthorizationContext + ) { + return delegated.workspaceId === canonicalContext.workspaceId + }, +} as const diff --git a/apps/sim/lib/knowledge/application/authorized-knowledge-use-case.ts b/apps/sim/lib/knowledge/application/authorized-knowledge-use-case.ts new file mode 100644 index 00000000000..5a384ef3647 --- /dev/null +++ b/apps/sim/lib/knowledge/application/authorized-knowledge-use-case.ts @@ -0,0 +1,28 @@ +import { + type AuthorizedWorkspaceUseCaseDefinition, + defineAuthorizedWorkspaceUseCase, + type WorkspaceOperation, +} from '@/lib/core/application' +import { + type KnowledgeAuthorizationContext, + knowledgeDelegationPolicy, +} from '@/lib/knowledge/application/authorization' + +type AuthorizedKnowledgeUseCaseDefinition< + O extends WorkspaceOperation, + I, + C extends KnowledgeAuthorizationContext, + R, +> = Omit, 'authorizationOptions'> + +export function defineAuthorizedKnowledgeUseCase< + const O extends WorkspaceOperation, + I, + C extends KnowledgeAuthorizationContext, + R, +>(definition: AuthorizedKnowledgeUseCaseDefinition) { + return defineAuthorizedWorkspaceUseCase({ + ...definition, + authorizationOptions: { delegation: knowledgeDelegationPolicy }, + }) +} diff --git a/apps/sim/lib/knowledge/application/billing.ts b/apps/sim/lib/knowledge/application/billing.ts new file mode 100644 index 00000000000..f2a08cd4f63 --- /dev/null +++ b/apps/sim/lib/knowledge/application/billing.ts @@ -0,0 +1,37 @@ +import type { Principal } from '@sim/auth/principal' +import { resolvePrincipalAttribution } from '@sim/auth/principal' +import { + type BillingAttributionSnapshot, + resolveBillingAttribution, + resolveSystemBillingAttribution, +} from '@/lib/billing/core/billing-attribution' +import type { KnowledgeWorkspaceContext } from '@/lib/knowledge/application/contexts' + +export class KnowledgeUsageLimitExceededError extends Error { + constructor(message: string) { + super(message) + this.name = 'KnowledgeUsageLimitExceededError' + } +} + +export function resolveKnowledgeAttributedUserId( + principal: Principal, + context: KnowledgeWorkspaceContext +): string { + return resolvePrincipalAttribution(principal, { + workspaceBillingOwnerUserId: context.billedAccountUserId, + }).attributedUserId +} + +export function resolveKnowledgeBillingAttribution( + principal: Principal, + context: KnowledgeWorkspaceContext +): Promise { + if (principal.kind === 'workspace_api_key') { + return resolveSystemBillingAttribution(context.workspaceId) + } + return resolveBillingAttribution({ + actorUserId: resolveKnowledgeAttributedUserId(principal, context), + workspaceId: context.workspaceId, + }) +} diff --git a/apps/sim/lib/knowledge/application/contexts.ts b/apps/sim/lib/knowledge/application/contexts.ts new file mode 100644 index 00000000000..1c2648aabd1 --- /dev/null +++ b/apps/sim/lib/knowledge/application/contexts.ts @@ -0,0 +1,83 @@ +import { db } from '@sim/db' +import { workspace } from '@sim/db/schema' +import { and, eq, isNull } from 'drizzle-orm' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import type { KnowledgeAuthorizationContext } from '@/lib/knowledge/application/authorization' +import type { ActiveKnowledgeDocument } from '@/lib/knowledge/documents/service' +import { getKnowledgeDocument } from '@/lib/knowledge/documents/service' +import { getKnowledgeBaseById } from '@/lib/knowledge/service' +import type { KnowledgeBaseWithCounts } from '@/lib/knowledge/types' + +export interface KnowledgeWorkspaceContext extends KnowledgeAuthorizationContext { + billedAccountUserId: string +} + +export interface ActiveKnowledgeBaseContext extends KnowledgeWorkspaceContext { + knowledgeBaseId: string + knowledgeBase: KnowledgeBaseWithCounts +} + +export interface ActiveKnowledgeDocumentContext extends ActiveKnowledgeBaseContext { + documentId: string + document: ActiveKnowledgeDocument +} + +export async function loadKnowledgeWorkspaceContext( + workspaceId: string +): Promise { + const [row] = await db + .select({ + workspaceId: workspace.id, + workspaceOrganizationId: workspace.organizationId, + allowPersonalApiKeys: workspace.allowPersonalApiKeys, + billedAccountUserId: workspace.billedAccountUserId, + }) + .from(workspace) + .where(and(eq(workspace.id, workspaceId), isNull(workspace.archivedAt))) + .limit(1) + return row ?? null +} + +export async function resolveKnowledgeWorkspaceContext(input: { + workspaceId: string +}): Promise { + const context = await loadKnowledgeWorkspaceContext(input.workspaceId) + if (!context) throw new OrchestrationError('not_found', 'Workspace not found') + return context +} + +export async function resolveActiveKnowledgeBaseContext(input: { + knowledgeBaseId: string + assertedWorkspaceId?: string +}): Promise { + const knowledgeBase = await getKnowledgeBaseById(input.knowledgeBaseId) + if ( + !knowledgeBase?.workspaceId || + (input.assertedWorkspaceId !== undefined && + knowledgeBase.workspaceId !== input.assertedWorkspaceId) + ) { + throw new OrchestrationError('not_found', 'Knowledge base not found') + } + const workspaceContext = await loadKnowledgeWorkspaceContext(knowledgeBase.workspaceId) + if (!workspaceContext) throw new OrchestrationError('not_found', 'Knowledge base not found') + return { + ...workspaceContext, + knowledgeBaseId: knowledgeBase.id, + knowledgeBase, + } +} + +export async function resolveActiveKnowledgeDocumentContext(input: { + knowledgeBaseId: string + documentId: string + assertedWorkspaceId?: string +}): Promise { + const context = await resolveActiveKnowledgeBaseContext(input) + const document = await getKnowledgeDocument(context.knowledgeBaseId, input.documentId) + if (!document) throw new OrchestrationError('not_found', 'Document not found') + return { + ...context, + documentId: document.id, + document, + } +} diff --git a/apps/sim/lib/knowledge/application/delegated-principal.ts b/apps/sim/lib/knowledge/application/delegated-principal.ts new file mode 100644 index 00000000000..ac0db06e992 --- /dev/null +++ b/apps/sim/lib/knowledge/application/delegated-principal.ts @@ -0,0 +1,33 @@ +import type { DelegatedPrincipal } from '@sim/auth/principal' +import { KNOWLEDGE_DELEGATION_AUDIENCE } from '@/lib/knowledge/application/authorization' + +const KNOWLEDGE_DELEGATION_TTL_MS = 5 * 60 * 1000 + +export interface CreateKnowledgeDelegatedPrincipalInput { + serviceId: DelegatedPrincipal['serviceId'] + subjectUserId: string + workspaceId: string + delegationId: string + chatId?: string + executionId?: string +} + +export function createKnowledgeDelegatedPrincipal( + input: CreateKnowledgeDelegatedPrincipalInput +): DelegatedPrincipal { + const issuedAt = new Date() + return { + kind: 'delegated', + serviceId: input.serviceId, + subjectUserId: input.subjectUserId, + workspaceId: input.workspaceId, + delegationId: input.delegationId, + audience: KNOWLEDGE_DELEGATION_AUDIENCE, + issuedAt, + expiresAt: new Date(issuedAt.getTime() + KNOWLEDGE_DELEGATION_TTL_MS), + resourceScope: { + ...(input.chatId ? { chatId: input.chatId } : {}), + ...(input.executionId ? { executionId: input.executionId } : {}), + }, + } +} diff --git a/apps/sim/lib/knowledge/application/documents.test.ts b/apps/sim/lib/knowledge/application/documents.test.ts new file mode 100644 index 00000000000..acfeea66d47 --- /dev/null +++ b/apps/sim/lib/knowledge/application/documents.test.ts @@ -0,0 +1,251 @@ +/** + * @vitest-environment node + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + resolveKnowledgeBase: vi.fn(), + resolveDocument: vi.fn(), + resolvePermission: vi.fn(), + resolveHumanBilling: vi.fn(), + resolveSystemBilling: vi.fn(), + checkUsage: vi.fn(), + getDocuments: vi.fn(), + createDocument: vi.fn(), + deleteDocument: vi.fn(), + processQueue: vi.fn(), + recordAudit: vi.fn(), +})) + +vi.mock('@sim/audit', () => ({ + AuditAction: { + DOCUMENT_UPLOADED: 'document.uploaded', + DOCUMENT_DELETED: 'document.deleted', + }, + AuditResourceType: { DOCUMENT: 'document' }, + recordAudit: mocks.recordAudit, +})) + +vi.mock('@sim/platform-authz/workspace', () => ({ + permissionSatisfies: (actual: string | null, required: string) => { + const rank = { read: 1, write: 2, admin: 3 } as const + return ( + actual !== null && rank[actual as keyof typeof rank] >= rank[required as keyof typeof rank] + ) + }, + resolveEffectiveWorkspacePermission: mocks.resolvePermission, +})) + +vi.mock('@/lib/billing/core/billing-attribution', () => ({ + resolveBillingAttribution: mocks.resolveHumanBilling, + resolveSystemBillingAttribution: mocks.resolveSystemBilling, + checkAttributedUsageLimits: mocks.checkUsage, +})) + +vi.mock('@/lib/knowledge/application/contexts', () => ({ + resolveActiveKnowledgeBaseContext: mocks.resolveKnowledgeBase, + resolveActiveKnowledgeDocumentContext: mocks.resolveDocument, +})) + +vi.mock('@/lib/knowledge/documents/service', () => ({ + getDocuments: mocks.getDocuments, + createSingleDocument: mocks.createDocument, + deleteKnowledgeDocumentInKnowledgeBase: mocks.deleteDocument, + processDocumentsWithQueue: mocks.processQueue, +})) + +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { + deleteKnowledgeDocument, + listKnowledgeDocuments, + uploadKnowledgeDocument, +} from '@/lib/knowledge/application/documents' + +const context = { + workspaceId: 'workspace-1', + workspaceOrganizationId: null, + allowPersonalApiKeys: true, + billedAccountUserId: 'billing-owner-1', + knowledgeBaseId: 'knowledge-1', + knowledgeBase: { id: 'knowledge-1', name: 'Docs' }, +} + +const document = { + id: 'document-1', + knowledgeBaseId: 'knowledge-1', + filename: 'guide.pdf', + fileUrl: '/api/files/serve/guide.pdf', + fileSize: 42, + mimeType: 'application/pdf', + enabled: true, + uploadedAt: new Date('2026-01-01T00:00:00Z'), +} + +describe('knowledge document application use cases', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.resolvePermission.mockResolvedValue('write') + mocks.resolveKnowledgeBase.mockResolvedValue(context) + mocks.resolveDocument.mockResolvedValue({ + ...context, + documentId: document.id, + document, + }) + mocks.resolveSystemBilling.mockResolvedValue({ + actorUserId: 'billing-owner-1', + workspaceId: 'workspace-1', + }) + mocks.resolveHumanBilling.mockResolvedValue({ + actorUserId: 'user-1', + workspaceId: 'workspace-1', + }) + mocks.checkUsage.mockResolvedValue({ isExceeded: false }) + mocks.createDocument.mockResolvedValue(document) + mocks.processQueue.mockResolvedValue(undefined) + mocks.getDocuments.mockResolvedValue({ + documents: [], + pagination: { total: 0, limit: 50, offset: 0, hasMore: false }, + }) + }) + + it('authorizes the canonical knowledge base before listing documents', async () => { + await listKnowledgeDocuments.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + limit: 25, + }, + }) + + expect(mocks.resolveKnowledgeBase).toHaveBeenCalledWith( + expect.objectContaining({ assertedWorkspaceId: 'workspace-1' }) + ) + expect(mocks.resolvePermission.mock.invocationCallOrder[0]).toBeLessThan( + mocks.getDocuments.mock.invocationCallOrder[0] + ) + }) + + it('resolves current workspace-key billing while retaining key audit attribution', async () => { + await uploadKnowledgeDocument.execute({ + principal: { + kind: 'workspace_api_key', + workspaceId: 'workspace-1', + keyId: 'key-1', + }, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + document, + source: 'v2', + }, + }) + + expect(mocks.resolveSystemBilling).toHaveBeenCalledWith('workspace-1') + expect(mocks.createDocument).toHaveBeenCalledWith( + document, + 'knowledge-1', + expect.any(String), + 'billing-owner-1', + undefined, + undefined, + { expectedWorkspaceId: 'workspace-1' } + ) + expect(mocks.recordAudit).toHaveBeenCalledWith( + expect.objectContaining({ + actorId: null, + actorName: 'Workspace API key', + metadata: expect.objectContaining({ + operation: 'knowledge.documents.upload', + actor: { + kind: 'workspace_api_key', + keyId: 'key-1', + workspaceId: 'workspace-1', + }, + }), + }) + ) + }) + + it('does not repeat usage admission after a code-defined pre-admission', async () => { + await uploadKnowledgeDocument.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + document, + usageAdmission: 'pre_admitted', + }, + }) + + expect(mocks.checkUsage).not.toHaveBeenCalled() + expect(mocks.createDocument).toHaveBeenCalledOnce() + }) + + it('conceals a cross-knowledge-base document before deletion and audit', async () => { + mocks.resolveDocument.mockRejectedValueOnce( + new OrchestrationError('not_found', 'Document not found') + ) + + await expect( + deleteKnowledgeDocument.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseId: 'knowledge-1', + documentId: 'document-from-another-kb', + assertedWorkspaceId: 'workspace-1', + }, + }) + ).rejects.toMatchObject({ code: 'not_found' }) + + expect(mocks.deleteDocument).not.toHaveBeenCalled() + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) + + it('carries canonical knowledge-base scope through deletion and audit', async () => { + await deleteKnowledgeDocument.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseId: 'knowledge-1', + documentId: 'document-1', + assertedWorkspaceId: 'workspace-1', + source: 'v2', + }, + }) + + expect(mocks.deleteDocument).toHaveBeenCalledWith( + 'knowledge-1', + 'document-1', + expect.any(String) + ) + expect(mocks.recordAudit).toHaveBeenCalledWith( + expect.objectContaining({ + action: 'document.deleted', + resourceId: 'document-1', + metadata: expect.objectContaining({ + operation: 'knowledge.documents.delete', + knowledgeBaseId: 'knowledge-1', + }), + }) + ) + }) + + it('propagates document infrastructure failures without audit', async () => { + const failure = new Error('storage ledger unavailable') + mocks.createDocument.mockRejectedValueOnce(failure) + + await expect( + uploadKnowledgeDocument.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + document, + }, + }) + ).rejects.toBe(failure) + + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/knowledge/application/documents.ts b/apps/sim/lib/knowledge/application/documents.ts new file mode 100644 index 00000000000..76f81c22e7d --- /dev/null +++ b/apps/sim/lib/knowledge/application/documents.ts @@ -0,0 +1,246 @@ +import { AuditAction, AuditResourceType } from '@sim/audit' +import { createLogger } from '@sim/logger' +import { checkAttributedUsageLimits } from '@/lib/billing/core/billing-attribution' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { generateRequestId } from '@/lib/core/utils/request' +import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' +import { + KnowledgeUsageLimitExceededError, + resolveKnowledgeAttributedUserId, + resolveKnowledgeBillingAttribution, +} from '@/lib/knowledge/application/billing' +import { + type ActiveKnowledgeDocumentContext, + resolveActiveKnowledgeBaseContext, + resolveActiveKnowledgeDocumentContext, +} from '@/lib/knowledge/application/contexts' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { + createSingleDocument, + type DocumentData, + deleteKnowledgeDocumentInKnowledgeBase, + getDocuments, + type ProcessingOptions, + processDocumentsWithQueue, +} from '@/lib/knowledge/documents/service' +import type { DocumentSortField, SortOrder } from '@/lib/knowledge/documents/types' +import { MAX_KNOWLEDGE_DOCUMENT_FILE_SIZE } from '@/lib/uploads/shared/types' +import { validateFileType } from '@/lib/uploads/utils/validation' + +const logger = createLogger('KnowledgeDocumentApplication') + +export interface ListKnowledgeDocumentsInput { + knowledgeBaseId: string + assertedWorkspaceId?: string + enabledFilter?: 'all' | 'enabled' | 'disabled' + search?: string + limit?: number + offset?: number + sortBy?: DocumentSortField + sortOrder?: SortOrder +} + +export interface ReadKnowledgeDocumentInput { + knowledgeBaseId: string + documentId: string + assertedWorkspaceId?: string +} + +export interface UploadKnowledgeDocumentAdmissionInput { + knowledgeBaseId: string + assertedWorkspaceId?: string +} + +export interface KnowledgeDocumentInput { + filename: string + fileUrl: string + fileSize: number + mimeType: string + documentTagsData?: string + tag1?: string + tag2?: string + tag3?: string + tag4?: string + tag5?: string + tag6?: string + tag7?: string +} + +export interface UploadKnowledgeDocumentInput extends UploadKnowledgeDocumentAdmissionInput { + document: KnowledgeDocumentInput + processingOptions?: ProcessingOptions + startProcessing?: boolean + /** Code-defined admission state; HTTP/model payloads must never populate it. */ + usageAdmission?: 'enforce' | 'pre_admitted' + source?: string +} + +export interface DeleteKnowledgeDocumentInput extends ReadKnowledgeDocumentInput { + source?: string +} + +export const listKnowledgeDocuments = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.listDocuments, + resolveContext: ({ input }: { input: ListKnowledgeDocumentsInput }) => + resolveActiveKnowledgeBaseContext(input), + async execute({ input, context }) { + const limit = input.limit ?? 50 + const offset = input.offset ?? 0 + if (!Number.isInteger(limit) || limit < 1 || limit > 100) { + throw new OrchestrationError('validation', 'Document limit must be between 1 and 100') + } + if (!Number.isInteger(offset) || offset < 0) { + throw new OrchestrationError('validation', 'Document offset must be a non-negative integer') + } + return getDocuments( + context.knowledgeBaseId, + { + enabledFilter: input.enabledFilter === 'all' ? undefined : input.enabledFilter, + search: input.search, + limit, + offset, + sortBy: input.sortBy, + sortOrder: input.sortOrder, + }, + generateRequestId() + ) + }, +}) + +export const readKnowledgeDocument = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.readDocument, + resolveContext: ({ input }: { input: ReadKnowledgeDocumentInput }) => + resolveActiveKnowledgeDocumentContext(input), + async execute({ context }: { context: ActiveKnowledgeDocumentContext }) { + return { document: context.document } + }, +}) + +export const admitKnowledgeDocumentUpload = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.uploadDocument, + resolveContext: ({ input }: { input: UploadKnowledgeDocumentAdmissionInput }) => + resolveActiveKnowledgeBaseContext(input), + async execute({ principal, context }) { + const billingAttribution = await resolveKnowledgeBillingAttribution(principal, context) + const usage = await checkAttributedUsageLimits(billingAttribution) + if (usage.isExceeded) { + throw new KnowledgeUsageLimitExceededError( + usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' + ) + } + return { + knowledgeBaseId: context.knowledgeBaseId, + knowledgeBaseName: context.knowledgeBase.name, + workspaceId: context.workspaceId, + storageActorUserId: resolveKnowledgeAttributedUserId(principal, context), + } + }, +}) + +export const uploadKnowledgeDocument = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.uploadDocument, + resolveContext: ({ input }: { input: UploadKnowledgeDocumentInput }) => + resolveActiveKnowledgeBaseContext(input), + async execute({ principal, input, context }) { + if (input.document.fileSize < 0 || input.document.fileSize > MAX_KNOWLEDGE_DOCUMENT_FILE_SIZE) { + throw new OrchestrationError( + 'payload_too_large', + 'Knowledge document exceeds the 100MB limit' + ) + } + const fileTypeError = validateFileType(input.document.filename, input.document.mimeType) + if (fileTypeError) throw new OrchestrationError('validation', fileTypeError.message) + const billingAttribution = await resolveKnowledgeBillingAttribution(principal, context) + if (input.usageAdmission !== 'pre_admitted') { + const usage = await checkAttributedUsageLimits(billingAttribution) + if (usage.isExceeded) { + throw new KnowledgeUsageLimitExceededError( + usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' + ) + } + } + const requestId = generateRequestId() + const uploadedBy = resolveKnowledgeAttributedUserId(principal, context) + const document = await createSingleDocument( + input.document, + context.knowledgeBaseId, + requestId, + uploadedBy, + undefined, + undefined, + { expectedWorkspaceId: context.workspaceId } + ) + if (input.startProcessing !== false) { + const processingDocument: DocumentData = { + documentId: document.id, + filename: document.filename, + fileUrl: document.fileUrl, + fileSize: document.fileSize, + mimeType: document.mimeType, + } + processDocumentsWithQueue( + [processingDocument], + context.knowledgeBaseId, + input.processingOptions ?? {}, + requestId, + billingAttribution + ).catch((error: unknown) => { + logger.error('Knowledge document processing pipeline failed', { + knowledgeBaseId: context.knowledgeBaseId, + documentId: document.id, + error, + }) + }) + } + return { document, created: true as const } + }, + projectAudit: ({ input, context, result }) => ({ + action: AuditAction.DOCUMENT_UPLOADED, + resourceType: AuditResourceType.DOCUMENT, + resourceId: result.document.id, + resourceName: result.document.filename, + description: `Uploaded document "${result.document.filename}" to knowledge base "${context.knowledgeBase.name}"`, + metadata: { + source: input.source, + knowledgeBaseId: context.knowledgeBaseId, + knowledgeBaseName: context.knowledgeBase.name, + fileName: result.document.filename, + fileType: result.document.mimeType, + fileSize: result.document.fileSize, + }, + }), +}) + +export const deleteKnowledgeDocument = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.deleteDocument, + resolveContext: ({ input }: { input: DeleteKnowledgeDocumentInput }) => + resolveActiveKnowledgeDocumentContext(input), + async execute({ context }: { context: ActiveKnowledgeDocumentContext }) { + await deleteKnowledgeDocumentInKnowledgeBase( + context.knowledgeBaseId, + context.documentId, + generateRequestId() + ) + return { + id: context.documentId, + filename: context.document.filename, + fileSize: context.document.fileSize, + mimeType: context.document.mimeType, + } + }, + projectAudit: ({ input, context, result }) => ({ + action: AuditAction.DOCUMENT_DELETED, + resourceType: AuditResourceType.DOCUMENT, + resourceId: result.id, + resourceName: result.filename, + description: `Deleted document "${result.filename}" from knowledge base "${context.knowledgeBase.name}"`, + metadata: { + source: input.source, + knowledgeBaseId: context.knowledgeBaseId, + knowledgeBaseName: context.knowledgeBase.name, + fileName: result.filename, + fileSize: result.fileSize, + mimeType: result.mimeType, + }, + }), +}) diff --git a/apps/sim/lib/knowledge/application/folder-paths.ts b/apps/sim/lib/knowledge/application/folder-paths.ts new file mode 100644 index 00000000000..cc1c9a70bf6 --- /dev/null +++ b/apps/sim/lib/knowledge/application/folder-paths.ts @@ -0,0 +1,33 @@ +import type { folder } from '@sim/db/schema' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { withFolderTreeLock } from '@/lib/folders/locks' +import type { FolderPathIndex } from '@/lib/folders/paths' +import { ROOT_FOLDER_PATH } from '@/lib/folders/paths' +import { loadActiveFolderPathIndex, resolveFolderPathFromIndex } from '@/lib/folders/queries' +import { MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE } from '@/lib/knowledge/constants' + +type FolderRow = typeof folder.$inferSelect + +export async function resolveKnowledgeFolderPath( + workspaceId: string, + path: string +): Promise<{ folderId: string | null; index: FolderPathIndex }> { + return withFolderTreeLock(workspaceId, 'knowledge_base', async (tx) => { + const index = await loadActiveFolderPathIndex(workspaceId, 'knowledge_base', tx, { + maxRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE, + }) + const folderId = resolveFolderPathFromIndex(index, path) + if (folderId === undefined) throw new OrchestrationError('not_found', 'Folder not found') + return { folderId, index } + }) +} + +export function knowledgeFolderPathForId( + index: FolderPathIndex, + folderId: string | null | undefined +): string { + if (!folderId) return ROOT_FOLDER_PATH + const path = index.pathById.get(folderId) + if (!path) throw new Error('Knowledge base references an inactive or missing folder') + return path +} diff --git a/apps/sim/lib/knowledge/application/folders.test.ts b/apps/sim/lib/knowledge/application/folders.test.ts new file mode 100644 index 00000000000..53e873224f7 --- /dev/null +++ b/apps/sim/lib/knowledge/application/folders.test.ts @@ -0,0 +1,195 @@ +/** + * @vitest-environment node + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + resolveWorkspace: vi.fn(), + resolvePermission: vi.fn(), + loadIndex: vi.fn(), + listRows: vi.fn(), + createAtPath: vi.fn(), + relocateByPath: vi.fn(), + deleteByPath: vi.fn(), + recordAudit: vi.fn(), + notify: vi.fn(), +})) + +vi.mock('@sim/audit', () => ({ + AuditAction: { + FOLDER_CREATED: 'folder.created', + FOLDER_MOVED: 'folder.moved', + FOLDER_DELETED: 'folder.deleted', + }, + AuditResourceType: { FOLDER: 'folder' }, + recordAudit: mocks.recordAudit, +})) + +vi.mock('@sim/platform-authz/workspace', () => ({ + permissionSatisfies: (actual: string | null, required: string) => { + const rank = { read: 1, write: 2, admin: 3 } as const + return ( + actual !== null && rank[actual as keyof typeof rank] >= rank[required as keyof typeof rank] + ) + }, + resolveEffectiveWorkspacePermission: mocks.resolvePermission, +})) + +vi.mock('@/lib/knowledge/application/contexts', () => ({ + resolveKnowledgeWorkspaceContext: mocks.resolveWorkspace, +})) + +vi.mock('@/lib/folders/queries', () => ({ + loadActiveFolderPathIndex: mocks.loadIndex, + listActiveFolderRows: mocks.listRows, + resolveFolderPathFromIndex: (index: { idByPath: Map }, path: string) => + path === '/' ? null : index.idByPath.get(path), +})) + +vi.mock('@/lib/folders/orchestration', () => ({ + createFolderAtPath: mocks.createAtPath, + relocateFolderByPath: mocks.relocateByPath, + deleteFolderByPath: mocks.deleteByPath, +})) + +vi.mock('@/lib/realtime/notify', () => ({ + notifyFolderResourceChanged: mocks.notify, +})) + +import { + createKnowledgeFolder, + deleteKnowledgeFolder, + listKnowledgeFolders, +} from '@/lib/knowledge/application/folders' + +const context = { + workspaceId: 'workspace-1', + workspaceOrganizationId: null, + allowPersonalApiKeys: true, + billedAccountUserId: 'billing-owner-1', +} + +const folder = { + id: 'folder-1', + resourceType: 'knowledge_base', + name: 'Docs', + userId: 'billing-owner-1', + workspaceId: 'workspace-1', + parentId: null, + sortOrder: 0, + locked: false, + createdAt: new Date('2026-01-01T00:00:00Z'), + updatedAt: new Date('2026-01-01T00:00:00Z'), + deletedAt: null, +} + +describe('knowledge folder application use cases', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.resolveWorkspace.mockResolvedValue(context) + mocks.resolvePermission.mockResolvedValue('write') + mocks.loadIndex.mockResolvedValue({ + idByPath: new Map([['/Docs', 'folder-1']]), + pathById: new Map([['folder-1', '/Docs']]), + rowById: new Map([['folder-1', folder]]), + }) + mocks.listRows.mockResolvedValue([folder]) + mocks.createAtPath.mockResolvedValue({ success: true, folder, path: '/Docs' }) + mocks.deleteByPath.mockResolvedValue({ + success: true, + path: '/Docs', + deletedItems: { folders: 2, knowledgeBases: 3 }, + }) + mocks.notify.mockResolvedValue(undefined) + }) + + it('resolves a canonical parent path before listing', async () => { + const result = await listKnowledgeFolders.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { workspaceId: 'workspace-1', parentPath: '/Docs' }, + }) + + expect(mocks.listRows).toHaveBeenCalledWith( + 'workspace-1', + 'knowledge_base', + expect.objectContaining({ parentId: 'folder-1' }) + ) + expect(result.folders[0]).toMatchObject({ id: 'folder-1', path: '/Docs' }) + }) + + it('rejects a missing parent without querying folder rows', async () => { + await expect( + listKnowledgeFolders.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { workspaceId: 'workspace-1', parentPath: '/Missing' }, + }) + ).rejects.toMatchObject({ code: 'not_found' }) + + expect(mocks.listRows).not.toHaveBeenCalled() + }) + + it('uses compatibility attribution only for storage and key attribution for audit', async () => { + await createKnowledgeFolder.execute({ + principal: { + kind: 'workspace_api_key', + workspaceId: 'workspace-1', + keyId: 'key-1', + }, + input: { workspaceId: 'workspace-1', path: '/Docs', source: 'v2' }, + }) + + expect(mocks.createAtPath).toHaveBeenCalledWith( + expect.objectContaining({ + resourceType: 'knowledge_base', + userId: 'billing-owner-1', + effects: false, + throwInfrastructure: true, + }) + ) + expect(mocks.recordAudit).toHaveBeenCalledWith( + expect.objectContaining({ + actorId: null, + actorName: 'Workspace API key', + metadata: expect.objectContaining({ + operation: 'knowledge.folders.create', + actor: { + kind: 'workspace_api_key', + keyId: 'key-1', + workspaceId: 'workspace-1', + }, + }), + }) + ) + expect(mocks.recordAudit.mock.invocationCallOrder[0]).toBeLessThan( + mocks.notify.mock.invocationCallOrder[0] + ) + }) + + it('preserves recursive cascade counts', async () => { + const result = await deleteKnowledgeFolder.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { workspaceId: 'workspace-1', path: '/Docs', recursive: true }, + }) + + expect(mocks.deleteByPath).toHaveBeenCalledWith( + expect.objectContaining({ recursive: true, effects: false, throwInfrastructure: true }) + ) + expect(result.deletedItems).toEqual({ folders: 2, knowledgeBases: 3 }) + }) + + it('propagates infrastructure failures without audit or notification', async () => { + const failure = new Error('folder database unavailable') + mocks.createAtPath.mockRejectedValueOnce(failure) + + await expect( + createKnowledgeFolder.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { workspaceId: 'workspace-1', path: '/Docs' }, + }) + ).rejects.toBe(failure) + + expect(mocks.recordAudit).not.toHaveBeenCalled() + expect(mocks.notify).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/knowledge/application/folders.ts b/apps/sim/lib/knowledge/application/folders.ts new file mode 100644 index 00000000000..5485322c039 --- /dev/null +++ b/apps/sim/lib/knowledge/application/folders.ts @@ -0,0 +1,209 @@ +import { AuditAction, AuditResourceType } from '@sim/audit' +import type { folder } from '@sim/db/schema' +import { OrchestrationError, type OrchestrationErrorCode } from '@/lib/core/orchestration/types' +import { + createFolderAtPath, + deleteFolderByPath, + relocateFolderByPath, +} from '@/lib/folders/orchestration' +import { ROOT_FOLDER_PATH } from '@/lib/folders/paths' +import { + listActiveFolderRows, + loadActiveFolderPathIndex, + resolveFolderPathFromIndex, +} from '@/lib/folders/queries' +import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' +import { resolveKnowledgeAttributedUserId } from '@/lib/knowledge/application/billing' +import { resolveKnowledgeWorkspaceContext } from '@/lib/knowledge/application/contexts' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE } from '@/lib/knowledge/constants' +import { notifyFolderResourceChanged } from '@/lib/realtime/notify' + +type KnowledgeFolder = typeof folder.$inferSelect & { path: string } + +export interface ListKnowledgeFoldersInput { + workspaceId: string + parentPath?: string + search?: string + sortBy?: 'name' | 'createdAt' | 'updatedAt' + sortOrder?: 'asc' | 'desc' +} + +export interface CreateKnowledgeFolderInput { + workspaceId: string + path: string + source?: string +} + +export interface RelocateKnowledgeFolderInput { + workspaceId: string + path: string + destinationPath: string + source?: string +} + +export interface DeleteKnowledgeFolderInput { + workspaceId: string + path: string + recursive?: boolean + source?: string +} + +function throwFolderFailure(result: { error?: string; errorCode?: OrchestrationErrorCode }): never { + throw new OrchestrationError( + result.errorCode ?? 'internal', + result.error ?? 'Folder operation failed' + ) +} + +export const listKnowledgeFolders = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.listFolders, + resolveContext: ({ input }: { input: ListKnowledgeFoldersInput }) => + resolveKnowledgeWorkspaceContext(input), + async execute({ input, context }) { + const index = await loadActiveFolderPathIndex( + context.workspaceId, + 'knowledge_base', + undefined, + { maxRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE } + ) + const parentId = + input.parentPath === undefined + ? undefined + : resolveFolderPathFromIndex(index, input.parentPath) + if (input.parentPath !== undefined && parentId === undefined) { + throw new OrchestrationError('not_found', 'Folder not found') + } + const folders = await listActiveFolderRows(context.workspaceId, 'knowledge_base', { + parentId, + search: input.search, + sortBy: input.sortBy, + sortOrder: input.sortOrder, + maxRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE, + }) + return { + folders: folders.map((folder): KnowledgeFolder => { + const path = index.pathById.get(folder.id) + if (!path) throw new Error('Folder path index is missing a listed folder') + return { ...folder, path } + }), + } + }, +}) + +export const createKnowledgeFolder = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.createFolder, + resolveContext: ({ input }: { input: CreateKnowledgeFolderInput }) => + resolveKnowledgeWorkspaceContext(input), + async execute({ principal, input, context }) { + const result = await createFolderAtPath({ + resourceType: 'knowledge_base', + workspaceId: context.workspaceId, + userId: resolveKnowledgeAttributedUserId(principal, context), + path: input.path, + effects: false, + throwInfrastructure: true, + maxFolderRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE, + }) + if (!result.success || !result.folder) return throwFolderFailure(result) + return { folder: { ...result.folder, path: result.path ?? input.path } } + }, + projectAudit: ({ input, result }) => ({ + action: AuditAction.FOLDER_CREATED, + resourceType: AuditResourceType.FOLDER, + resourceId: result.folder.id, + resourceName: result.folder.name, + description: `Created knowledge base folder "${result.folder.path}"`, + metadata: { + source: input.source, + path: result.folder.path, + folderResourceType: 'knowledge_base', + }, + }), + afterSuccess: ({ context }) => notifyFolderResourceChanged('knowledge_base', context.workspaceId), +}) + +export const relocateKnowledgeFolder = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.relocateFolder, + resolveContext: ({ input }: { input: RelocateKnowledgeFolderInput }) => + resolveKnowledgeWorkspaceContext(input), + async execute({ principal, input, context }) { + const result = await relocateFolderByPath({ + resourceType: 'knowledge_base', + workspaceId: context.workspaceId, + userId: resolveKnowledgeAttributedUserId(principal, context), + path: input.path, + destinationPath: input.destinationPath, + effects: false, + throwInfrastructure: true, + maxFolderRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE, + }) + if (!result.success || !result.folder) return throwFolderFailure(result) + return { folder: { ...result.folder, path: result.path ?? input.destinationPath } } + }, + projectAudit: ({ input, result }) => ({ + action: AuditAction.FOLDER_MOVED, + resourceType: AuditResourceType.FOLDER, + resourceId: result.folder.id, + resourceName: result.folder.name, + description: `Moved knowledge base folder to "${result.folder.path}"`, + metadata: { + source: input.source, + sourcePath: input.path, + destinationPath: result.folder.path, + folderResourceType: 'knowledge_base', + }, + }), + afterSuccess: ({ context }) => notifyFolderResourceChanged('knowledge_base', context.workspaceId), +}) + +export const deleteKnowledgeFolder = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.deleteFolder, + resolveContext: ({ input }: { input: DeleteKnowledgeFolderInput }) => + resolveKnowledgeWorkspaceContext(input), + async execute({ principal, input, context }) { + if (input.path === ROOT_FOLDER_PATH) { + throw new OrchestrationError('validation', 'Cannot delete the root path') + } + const index = await loadActiveFolderPathIndex( + context.workspaceId, + 'knowledge_base', + undefined, + { maxRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE } + ) + const folderId = resolveFolderPathFromIndex(index, input.path) + const folder = typeof folderId === 'string' ? index.rowById.get(folderId) : undefined + if (!folder) throw new OrchestrationError('not_found', 'Folder not found') + const result = await deleteFolderByPath({ + resourceType: 'knowledge_base', + workspaceId: context.workspaceId, + userId: resolveKnowledgeAttributedUserId(principal, context), + path: input.path, + recursive: input.recursive ?? false, + effects: false, + throwInfrastructure: true, + maxFolderRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE, + }) + if (!result.success || !result.deletedItems) return throwFolderFailure(result) + return { + id: folder.id, + name: folder.name, + path: input.path, + deletedItems: result.deletedItems, + } + }, + projectAudit: ({ input, result }) => ({ + action: AuditAction.FOLDER_DELETED, + resourceType: AuditResourceType.FOLDER, + resourceId: result.id, + resourceName: result.name, + description: `Deleted knowledge base folder "${result.path}"`, + metadata: { + source: input.source, + path: result.path, + folderResourceType: 'knowledge_base', + deletedItems: result.deletedItems, + }, + }), + afterSuccess: ({ context }) => notifyFolderResourceChanged('knowledge_base', context.workspaceId), +}) diff --git a/apps/sim/lib/knowledge/application/knowledge-bases.test.ts b/apps/sim/lib/knowledge/application/knowledge-bases.test.ts new file mode 100644 index 00000000000..9ae6b3f7fc9 --- /dev/null +++ b/apps/sim/lib/knowledge/application/knowledge-bases.test.ts @@ -0,0 +1,209 @@ +/** + * @vitest-environment node + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + resolveWorkspace: vi.fn(), + resolveKnowledgeBase: vi.fn(), + resolvePermission: vi.fn(), + resolveFolderPath: vi.fn(), + createRecord: vi.fn(), + updateRecord: vi.fn(), + deleteRecord: vi.fn(), + listRecords: vi.fn(), + loadFolderIndex: vi.fn(), + recordAudit: vi.fn(), +})) + +vi.mock('@sim/audit', () => ({ + AuditAction: { + KNOWLEDGE_BASE_CREATED: 'knowledge_base.created', + KNOWLEDGE_BASE_UPDATED: 'knowledge_base.updated', + KNOWLEDGE_BASE_DELETED: 'knowledge_base.deleted', + }, + AuditResourceType: { KNOWLEDGE_BASE: 'knowledge_base' }, + recordAudit: mocks.recordAudit, +})) + +vi.mock('@sim/platform-authz/workspace', () => ({ + permissionSatisfies: (actual: string | null, required: string) => { + const rank = { read: 1, write: 2, admin: 3 } as const + return ( + actual !== null && rank[actual as keyof typeof rank] >= rank[required as keyof typeof rank] + ) + }, + resolveEffectiveWorkspacePermission: mocks.resolvePermission, +})) + +vi.mock('@/lib/folders/queries', () => ({ + loadActiveFolderPathIndex: mocks.loadFolderIndex, +})) + +vi.mock('@/lib/knowledge/application/contexts', () => ({ + resolveKnowledgeWorkspaceContext: mocks.resolveWorkspace, + resolveActiveKnowledgeBaseContext: mocks.resolveKnowledgeBase, +})) + +vi.mock('@/lib/knowledge/application/folder-paths', () => ({ + resolveKnowledgeFolderPath: mocks.resolveFolderPath, + knowledgeFolderPathForId: () => '/', +})) + +vi.mock('@/lib/knowledge/embeddings', () => ({ + EMBEDDING_DIMENSIONS: 1536, + getConfiguredEmbeddingModel: () => 'text-embedding-3-small', +})) + +vi.mock('@/lib/knowledge/service', () => ({ + createAuthorizedKnowledgeBase: mocks.createRecord, + updateKnowledgeBase: mocks.updateRecord, + deleteKnowledgeBase: mocks.deleteRecord, + getWorkspaceKnowledgeBases: mocks.listRecords, +})) + +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { + createKnowledgeBase, + readKnowledgeBase, + updateKnowledgeBaseOperation, +} from '@/lib/knowledge/application/knowledge-bases' + +const context = { + workspaceId: 'workspace-1', + workspaceOrganizationId: 'organization-1', + allowPersonalApiKeys: true, + billedAccountUserId: 'billing-owner-1', +} + +const knowledgeBase = { + id: 'knowledge-1', + userId: 'billing-owner-1', + name: 'Docs', + description: null, + tokenCount: 0, + embeddingModel: 'text-embedding-3-small', + embeddingDimension: 1536, + chunkingConfig: { maxSize: 1024, minSize: 100, overlap: 200 }, + createdAt: new Date('2026-01-01T00:00:00Z'), + updatedAt: new Date('2026-01-01T00:00:00Z'), + deletedAt: null, + workspaceId: 'workspace-1', + folderId: null, + docCount: 0, + connectorTypes: [], +} + +describe('knowledge base application use cases', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.resolveWorkspace.mockResolvedValue(context) + mocks.resolveKnowledgeBase.mockResolvedValue({ + ...context, + knowledgeBaseId: knowledgeBase.id, + knowledgeBase, + }) + mocks.resolvePermission.mockResolvedValue('write') + mocks.resolveFolderPath.mockResolvedValue({ + folderId: null, + index: { pathById: new Map(), idByPath: new Map(), rowById: new Map() }, + }) + mocks.loadFolderIndex.mockResolvedValue({ pathById: new Map() }) + mocks.createRecord.mockResolvedValue(knowledgeBase) + mocks.updateRecord.mockResolvedValue({ ...knowledgeBase, name: 'Renamed' }) + }) + + it('rejects an insufficient role before the protected mutation', async () => { + mocks.resolvePermission.mockResolvedValueOnce('read') + + await expect( + createKnowledgeBase.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { workspaceId: 'workspace-1', name: 'Docs' }, + }) + ).rejects.toMatchObject({ code: 'forbidden' }) + + expect(mocks.createRecord).not.toHaveBeenCalled() + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) + + it('uses billing ownership only for the workspace-key compatibility column', async () => { + await createKnowledgeBase.execute({ + principal: { + kind: 'workspace_api_key', + workspaceId: 'workspace-1', + keyId: 'workspace-key-1', + }, + input: { workspaceId: 'workspace-1', name: 'Docs', source: 'v2' }, + }) + + expect(mocks.resolvePermission).not.toHaveBeenCalled() + expect(mocks.createRecord).toHaveBeenCalledWith( + expect.objectContaining({ userId: 'billing-owner-1', workspaceId: 'workspace-1' }), + expect.any(String) + ) + expect(mocks.recordAudit).toHaveBeenCalledWith( + expect.objectContaining({ + actorId: null, + actorName: 'Workspace API key', + metadata: expect.objectContaining({ + operation: 'knowledge.create', + actor: { + kind: 'workspace_api_key', + keyId: 'workspace-key-1', + workspaceId: 'workspace-1', + }, + }), + }) + ) + }) + + it('conceals a canonical scope mismatch and never audits it', async () => { + mocks.resolveKnowledgeBase.mockRejectedValueOnce( + new OrchestrationError('not_found', 'Knowledge base not found') + ) + + await expect( + readKnowledgeBase.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { knowledgeBaseId: 'knowledge-1', assertedWorkspaceId: 'workspace-2' }, + }) + ).rejects.toMatchObject({ code: 'not_found' }) + + expect(mocks.resolvePermission).not.toHaveBeenCalled() + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) + + it('propagates infrastructure failures without audit', async () => { + const failure = new Error('database unavailable') + mocks.createRecord.mockRejectedValueOnce(failure) + + await expect( + createKnowledgeBase.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { workspaceId: 'workspace-1', name: 'Docs' }, + }) + ).rejects.toBe(failure) + + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) + + it('carries the canonical workspace predicate into the locked update', async () => { + await updateKnowledgeBaseOperation.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + name: 'Renamed', + }, + }) + + expect(mocks.updateRecord).toHaveBeenCalledWith( + 'knowledge-1', + expect.objectContaining({ name: 'Renamed' }), + expect.any(String), + { assertedWorkspaceId: 'workspace-1' } + ) + }) +}) diff --git a/apps/sim/lib/knowledge/application/knowledge-bases.ts b/apps/sim/lib/knowledge/application/knowledge-bases.ts new file mode 100644 index 00000000000..56d4ab1c2db --- /dev/null +++ b/apps/sim/lib/knowledge/application/knowledge-bases.ts @@ -0,0 +1,268 @@ +import { AuditAction, AuditResourceType } from '@sim/audit' +import { createLogger } from '@sim/logger' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { generateRequestId } from '@/lib/core/utils/request' +import { loadActiveFolderPathIndex } from '@/lib/folders/queries' +import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' +import { resolveKnowledgeAttributedUserId } from '@/lib/knowledge/application/billing' +import { + type ActiveKnowledgeBaseContext, + type KnowledgeWorkspaceContext, + resolveActiveKnowledgeBaseContext, + resolveKnowledgeWorkspaceContext, +} from '@/lib/knowledge/application/contexts' +import { + knowledgeFolderPathForId, + resolveKnowledgeFolderPath, +} from '@/lib/knowledge/application/folder-paths' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { + DEFAULT_CHUNKING_CONFIG, + MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE, +} from '@/lib/knowledge/constants' +import { EMBEDDING_DIMENSIONS, getConfiguredEmbeddingModel } from '@/lib/knowledge/embeddings' +import { + createAuthorizedKnowledgeBase, + deleteKnowledgeBase, + getWorkspaceKnowledgeBases, + updateKnowledgeBase, +} from '@/lib/knowledge/service' +import type { ChunkingConfig, KnowledgeBaseWithCounts } from '@/lib/knowledge/types' + +const logger = createLogger('KnowledgeBaseApplication') + +export interface ListKnowledgeBasesInput { + workspaceId: string + folderPath?: string + search?: string + sortBy?: 'name' | 'createdAt' | 'updatedAt' + sortOrder?: 'asc' | 'desc' +} + +export interface KnowledgeBaseResult { + knowledgeBase: KnowledgeBaseWithCounts + folderPath: string +} + +export interface ListKnowledgeBasesResult { + knowledgeBases: KnowledgeBaseResult[] +} + +export interface CreateKnowledgeBaseInput { + workspaceId: string + name: string + description?: string + chunkingConfig?: Partial + folderPath?: string + source?: string +} + +export interface ReadKnowledgeBaseInput { + knowledgeBaseId: string + assertedWorkspaceId?: string +} + +export interface UpdateKnowledgeBaseInput extends ReadKnowledgeBaseInput { + name?: string + description?: string + chunkingConfig?: ChunkingConfig + folderPath?: string + source?: string +} + +export interface DeleteKnowledgeBaseInput extends ReadKnowledgeBaseInput { + source?: string +} + +async function executeListKnowledgeBases(args: { + input: ListKnowledgeBasesInput + context: KnowledgeWorkspaceContext +}): Promise { + const index = await loadActiveFolderPathIndex( + args.context.workspaceId, + 'knowledge_base', + undefined, + { maxRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE } + ) + const folderId = + args.input.folderPath === undefined + ? undefined + : await resolveKnowledgeFolderPath(args.context.workspaceId, args.input.folderPath).then( + (resolved) => resolved.folderId + ) + const rows = await getWorkspaceKnowledgeBases(args.context.workspaceId, 'active', { + folderId, + search: args.input.search, + sortBy: args.input.sortBy, + sortOrder: args.input.sortOrder, + }) + return { + knowledgeBases: rows.map((knowledgeBase) => ({ + knowledgeBase, + folderPath: knowledgeFolderPathForId(index, knowledgeBase.folderId), + })), + } +} + +async function executeCreateKnowledgeBase(args: { + principal: Parameters[0] + input: CreateKnowledgeBaseInput + context: KnowledgeWorkspaceContext +}): Promise { + const path = args.input.folderPath ?? '/' + const { folderId, index } = await resolveKnowledgeFolderPath(args.context.workspaceId, path) + const chunkingConfig: ChunkingConfig = { + ...DEFAULT_CHUNKING_CONFIG, + ...args.input.chunkingConfig, + } + const knowledgeBase = await createAuthorizedKnowledgeBase( + { + name: args.input.name, + description: args.input.description, + workspaceId: args.context.workspaceId, + folderId, + userId: resolveKnowledgeAttributedUserId(args.principal, args.context), + embeddingModel: getConfiguredEmbeddingModel(), + embeddingDimension: EMBEDDING_DIMENSIONS, + chunkingConfig, + }, + generateRequestId() + ) + logger.info('Created knowledge base', { + workspaceId: args.context.workspaceId, + knowledgeBaseId: knowledgeBase.id, + principalKind: args.principal.kind, + }) + return { knowledgeBase, folderPath: knowledgeFolderPathForId(index, knowledgeBase.folderId) } +} + +async function executeReadKnowledgeBase(args: { + context: ActiveKnowledgeBaseContext +}): Promise { + const index = await loadActiveFolderPathIndex( + args.context.workspaceId, + 'knowledge_base', + undefined, + { maxRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE } + ) + return { + knowledgeBase: args.context.knowledgeBase, + folderPath: knowledgeFolderPathForId(index, args.context.knowledgeBase.folderId), + } +} + +async function executeUpdateKnowledgeBase(args: { + input: UpdateKnowledgeBaseInput + context: ActiveKnowledgeBaseContext +}): Promise { + const updates = { + name: args.input.name, + description: args.input.description, + chunkingConfig: args.input.chunkingConfig, + folderId: + args.input.folderPath === undefined + ? undefined + : (await resolveKnowledgeFolderPath(args.context.workspaceId, args.input.folderPath)) + .folderId, + } + if (Object.values(updates).every((value) => value === undefined)) { + throw new OrchestrationError('validation', 'No updates specified') + } + const knowledgeBase = await updateKnowledgeBase( + args.context.knowledgeBaseId, + updates, + generateRequestId(), + { assertedWorkspaceId: args.context.workspaceId } + ) + const index = await loadActiveFolderPathIndex( + args.context.workspaceId, + 'knowledge_base', + undefined, + { maxRows: MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE } + ) + logger.info('Updated knowledge base', { + workspaceId: args.context.workspaceId, + knowledgeBaseId: knowledgeBase.id, + }) + return { knowledgeBase, folderPath: knowledgeFolderPathForId(index, knowledgeBase.folderId) } +} + +async function executeDeleteKnowledgeBase(args: { + context: ActiveKnowledgeBaseContext +}): Promise<{ id: string; name: string }> { + await deleteKnowledgeBase(args.context.knowledgeBaseId, generateRequestId(), { + assertedWorkspaceId: args.context.workspaceId, + }) + return { id: args.context.knowledgeBaseId, name: args.context.knowledgeBase.name } +} + +export const listKnowledgeBases = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.list, + resolveContext: ({ input }: { input: ListKnowledgeBasesInput }) => + resolveKnowledgeWorkspaceContext(input), + execute: executeListKnowledgeBases, +}) + +export const createKnowledgeBase = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.create, + resolveContext: ({ input }: { input: CreateKnowledgeBaseInput }) => + resolveKnowledgeWorkspaceContext(input), + execute: executeCreateKnowledgeBase, + projectAudit: ({ input, result }) => ({ + action: AuditAction.KNOWLEDGE_BASE_CREATED, + resourceType: AuditResourceType.KNOWLEDGE_BASE, + resourceId: result.knowledgeBase.id, + resourceName: result.knowledgeBase.name, + description: `Created knowledge base "${result.knowledgeBase.name}"`, + metadata: { + source: input.source, + name: result.knowledgeBase.name, + description: result.knowledgeBase.description, + embeddingModel: result.knowledgeBase.embeddingModel, + embeddingDimension: result.knowledgeBase.embeddingDimension, + folderPath: result.folderPath, + }, + }), +}) + +export const readKnowledgeBase = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.read, + resolveContext: ({ input }: { input: ReadKnowledgeBaseInput }) => + resolveActiveKnowledgeBaseContext(input), + execute: executeReadKnowledgeBase, +}) + +export const updateKnowledgeBaseOperation = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.update, + resolveContext: ({ input }: { input: UpdateKnowledgeBaseInput }) => + resolveActiveKnowledgeBaseContext(input), + execute: executeUpdateKnowledgeBase, + projectAudit: ({ input, result }) => ({ + action: AuditAction.KNOWLEDGE_BASE_UPDATED, + resourceType: AuditResourceType.KNOWLEDGE_BASE, + resourceId: result.knowledgeBase.id, + resourceName: result.knowledgeBase.name, + description: `Updated knowledge base "${result.knowledgeBase.name}"`, + metadata: { + source: input.source, + updatedFields: ['name', 'description', 'chunkingConfig', 'folderPath'].filter( + (key) => input[key as keyof UpdateKnowledgeBaseInput] !== undefined + ), + }, + }), +}) + +export const deleteKnowledgeBaseOperation = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.delete, + resolveContext: ({ input }: { input: DeleteKnowledgeBaseInput }) => + resolveActiveKnowledgeBaseContext(input), + execute: executeDeleteKnowledgeBase, + projectAudit: ({ input, result }) => ({ + action: AuditAction.KNOWLEDGE_BASE_DELETED, + resourceType: AuditResourceType.KNOWLEDGE_BASE, + resourceId: result.id, + resourceName: result.name, + description: `Deleted knowledge base "${result.name}"`, + metadata: { source: input.source, knowledgeBaseName: result.name }, + }), +}) diff --git a/apps/sim/lib/knowledge/application/operations.test.ts b/apps/sim/lib/knowledge/application/operations.test.ts new file mode 100644 index 00000000000..49b12c42476 --- /dev/null +++ b/apps/sim/lib/knowledge/application/operations.test.ts @@ -0,0 +1,50 @@ +/** + * @vitest-environment node + */ + +import { permissionSatisfies } from '@sim/platform-authz/workspace' +import { describe, expect, it } from 'vitest' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' + +describe('knowledge operation registry', () => { + it('defines unique stable semantic operation IDs', () => { + const ids = Object.values(knowledgeOperations).map((operation) => operation.id) + expect(ids).toEqual([ + 'knowledge.list', + 'knowledge.read', + 'knowledge.create', + 'knowledge.update', + 'knowledge.delete', + 'knowledge.search', + 'knowledge.folders.list', + 'knowledge.folders.create', + 'knowledge.folders.relocate', + 'knowledge.folders.delete', + 'knowledge.documents.list', + 'knowledge.documents.read', + 'knowledge.documents.upload', + 'knowledge.documents.delete', + 'knowledge.documents.upload.create', + 'knowledge.documents.upload.parts', + 'knowledge.documents.upload.complete', + 'knowledge.documents.upload.cancel', + ]) + expect(new Set(ids).size).toBe(ids.length) + }) + + it('keeps workspace keys within their fixed write ceiling', () => { + for (const operation of Object.values(knowledgeOperations)) { + expect(operation.workspaceApiKey).toBe('allow') + expect(operation.principalKinds).toContain('workspace_api_key') + expect(permissionSatisfies('write', operation.minimumRole)).toBe(true) + } + }) + + it('allows delegated callers only on semantic knowledge and document operations', () => { + expect(knowledgeOperations.list.principalKinds).toContain('delegated') + expect(knowledgeOperations.search.principalKinds).toContain('delegated') + expect(knowledgeOperations.uploadDocument.principalKinds).toContain('delegated') + expect(knowledgeOperations.listFolders.principalKinds).not.toContain('delegated') + expect(knowledgeOperations.uploadComplete.principalKinds).not.toContain('delegated') + }) +}) diff --git a/apps/sim/lib/knowledge/application/operations.ts b/apps/sim/lib/knowledge/application/operations.ts new file mode 100644 index 00000000000..199ab6931d1 --- /dev/null +++ b/apps/sim/lib/knowledge/application/operations.ts @@ -0,0 +1,123 @@ +import { defineWorkspaceOperation } from '@/lib/core/application' + +const ALL_PRINCIPAL_KINDS = [ + 'session', + 'personal_api_key', + 'workspace_api_key', + 'delegated', +] as const + +const HTTP_PRINCIPAL_KINDS = ['session', 'personal_api_key', 'workspace_api_key'] as const + +export const knowledgeOperations = { + list: defineWorkspaceOperation({ + id: 'knowledge.list', + minimumRole: 'read', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + read: defineWorkspaceOperation({ + id: 'knowledge.read', + minimumRole: 'read', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + create: defineWorkspaceOperation({ + id: 'knowledge.create', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + update: defineWorkspaceOperation({ + id: 'knowledge.update', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + delete: defineWorkspaceOperation({ + id: 'knowledge.delete', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + search: defineWorkspaceOperation({ + id: 'knowledge.search', + minimumRole: 'read', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + listFolders: defineWorkspaceOperation({ + id: 'knowledge.folders.list', + minimumRole: 'read', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), + createFolder: defineWorkspaceOperation({ + id: 'knowledge.folders.create', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), + relocateFolder: defineWorkspaceOperation({ + id: 'knowledge.folders.relocate', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), + deleteFolder: defineWorkspaceOperation({ + id: 'knowledge.folders.delete', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), + listDocuments: defineWorkspaceOperation({ + id: 'knowledge.documents.list', + minimumRole: 'read', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + readDocument: defineWorkspaceOperation({ + id: 'knowledge.documents.read', + minimumRole: 'read', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + uploadDocument: defineWorkspaceOperation({ + id: 'knowledge.documents.upload', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + deleteDocument: defineWorkspaceOperation({ + id: 'knowledge.documents.delete', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: ALL_PRINCIPAL_KINDS, + }), + uploadCreate: defineWorkspaceOperation({ + id: 'knowledge.documents.upload.create', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), + uploadParts: defineWorkspaceOperation({ + id: 'knowledge.documents.upload.parts', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), + uploadComplete: defineWorkspaceOperation({ + id: 'knowledge.documents.upload.complete', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), + uploadCancel: defineWorkspaceOperation({ + id: 'knowledge.documents.upload.cancel', + minimumRole: 'write', + workspaceApiKey: 'allow', + principalKinds: HTTP_PRINCIPAL_KINDS, + }), +} as const + +export type KnowledgeOperation = (typeof knowledgeOperations)[keyof typeof knowledgeOperations] diff --git a/apps/sim/lib/knowledge/application/search.test.ts b/apps/sim/lib/knowledge/application/search.test.ts new file mode 100644 index 00000000000..a56e6350ac4 --- /dev/null +++ b/apps/sim/lib/knowledge/application/search.test.ts @@ -0,0 +1,272 @@ +/** + * @vitest-environment node + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + resolveWorkspace: vi.fn(), + resolvePermission: vi.fn(), + getKnowledgeBase: vi.fn(), + resolveBilling: vi.fn(), + checkUsage: vi.fn(), + generateEmbedding: vi.fn(), + executeSearch: vi.fn(), + getDocumentMetadata: vi.fn(), + getTagDefinitions: vi.fn(), + recordEmbeddingUsage: vi.fn(), + importProvenance: vi.fn(), +})) + +vi.mock('@sim/platform-authz/workspace', () => ({ + permissionSatisfies: (actual: string | null, required: string) => { + const rank = { read: 1, write: 2, admin: 3 } as const + return ( + actual !== null && rank[actual as keyof typeof rank] >= rank[required as keyof typeof rank] + ) + }, + resolveEffectiveWorkspacePermission: mocks.resolvePermission, +})) + +vi.mock('@/lib/billing/core/billing-attribution', () => ({ + resolveBillingAttribution: mocks.resolveBilling, + resolveSystemBillingAttribution: mocks.resolveBilling, + checkAttributedUsageLimits: mocks.checkUsage, +})) + +vi.mock('@/lib/knowledge/application/contexts', () => ({ + resolveKnowledgeWorkspaceContext: mocks.resolveWorkspace, +})) + +vi.mock('@/lib/knowledge/service', () => ({ + getKnowledgeBaseById: mocks.getKnowledgeBase, +})) + +vi.mock('@/lib/knowledge/embeddings', () => ({ + generateSearchEmbedding: mocks.generateEmbedding, + recordSearchEmbeddingUsage: mocks.recordEmbeddingUsage, +})) + +vi.mock('@/lib/knowledge/search/queries', () => ({ + generateSearchEmbedding: mocks.generateEmbedding, + executeKnowledgeSearch: mocks.executeSearch, + getDocumentMetadataByIds: mocks.getDocumentMetadata, +})) + +vi.mock('@/lib/knowledge/tags/service', () => ({ + getDocumentTagDefinitions: mocks.getTagDefinitions, +})) + +vi.mock('@/lib/knowledge/tags/utils', () => ({ + buildUndefinedTagsError: (tags: string[]) => `Undefined tags: ${tags.join(', ')}`, + validateTagValue: () => null, +})) + +vi.mock('@/lib/knowledge/secret-provenance', () => ({ + importKnowledgeSearchResultSecretProvenance: mocks.importProvenance, +})) + +import { searchKnowledge } from '@/lib/knowledge/application/search' + +const workspace = { + workspaceId: 'workspace-1', + workspaceOrganizationId: null, + allowPersonalApiKeys: true, + billedAccountUserId: 'billing-owner-1', +} + +const knowledgeBase = { + id: 'knowledge-1', + workspaceId: 'workspace-1', + embeddingModel: 'text-embedding-3-small', +} + +describe('knowledge search application use case', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.resolveWorkspace.mockResolvedValue(workspace) + mocks.resolvePermission.mockResolvedValue('read') + mocks.getKnowledgeBase.mockResolvedValue(knowledgeBase) + mocks.resolveBilling.mockResolvedValue({ + actorUserId: 'user-1', + workspaceId: 'workspace-1', + }) + mocks.checkUsage.mockResolvedValue({ isExceeded: false }) + mocks.generateEmbedding.mockResolvedValue({ embedding: [0.1], isBYOK: false }) + mocks.executeSearch.mockResolvedValue([ + { + id: 'embedding-1', + documentId: 'document-1', + knowledgeBaseId: 'knowledge-1', + content: 'answer', + chunkIndex: 0, + distance: 0.2, + tag1: null, + tag2: null, + tag3: null, + tag4: null, + tag5: null, + tag6: null, + tag7: null, + number1: null, + number2: null, + number3: null, + number4: null, + number5: null, + date1: null, + date2: null, + boolean1: null, + boolean2: null, + boolean3: null, + }, + ]) + mocks.getDocumentMetadata.mockResolvedValue({ + 'document-1': { filename: 'guide.pdf', sourceUrl: null }, + }) + mocks.getTagDefinitions.mockResolvedValue([]) + mocks.recordEmbeddingUsage.mockResolvedValue(undefined) + mocks.importProvenance.mockResolvedValue({ imported: true, documentMetadata: {} }) + }) + + it('authorizes every canonical knowledge base before billing and search', async () => { + const result = await searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + workspaceId: 'workspace-1', + knowledgeBaseIds: ['knowledge-1'], + query: 'answer', + topK: 5, + }, + }) + + expect(mocks.resolvePermission.mock.invocationCallOrder[0]).toBeLessThan( + mocks.resolveBilling.mock.invocationCallOrder[0] + ) + expect(mocks.resolveBilling.mock.invocationCallOrder[0]).toBeLessThan( + mocks.executeSearch.mock.invocationCallOrder[0] + ) + expect(mocks.executeSearch).toHaveBeenCalledWith( + expect.objectContaining({ + knowledgeBaseIds: ['knowledge-1'], + topK: 5, + searchMode: 'vector', + }) + ) + expect(result.results[0]).toMatchObject({ + embeddingId: 'embedding-1', + documentId: 'document-1', + similarity: 0.8, + }) + }) + + it('rejects a cross-workspace knowledge base before authorization or spend', async () => { + mocks.getKnowledgeBase.mockResolvedValueOnce({ + ...knowledgeBase, + workspaceId: 'workspace-2', + }) + + await expect( + searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + workspaceId: 'workspace-1', + knowledgeBaseIds: ['knowledge-1'], + query: 'answer', + topK: 5, + }, + }) + ).rejects.toMatchObject({ code: 'not_found' }) + + expect(mocks.resolvePermission).not.toHaveBeenCalled() + expect(mocks.resolveBilling).not.toHaveBeenCalled() + expect(mocks.executeSearch).not.toHaveBeenCalled() + }) + + it('enforces semantic knowledge-base and result bounds for trusted callers', async () => { + await expect( + searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + workspaceId: 'workspace-1', + knowledgeBaseIds: Array.from({ length: 21 }, (_, index) => `knowledge-${index}`), + query: 'answer', + topK: 5, + }, + }) + ).rejects.toMatchObject({ code: 'validation' }) + + await expect( + searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + workspaceId: 'workspace-1', + knowledgeBaseIds: ['knowledge-1'], + query: 'answer', + topK: 101, + }, + }) + ).rejects.toMatchObject({ code: 'validation' }) + + expect(mocks.resolveWorkspace).not.toHaveBeenCalled() + expect(mocks.getKnowledgeBase).not.toHaveBeenCalled() + }) + + it('rejects multi-knowledge-base tag filters without embedding spend', async () => { + mocks.getKnowledgeBase + .mockResolvedValueOnce(knowledgeBase) + .mockResolvedValueOnce({ ...knowledgeBase, id: 'knowledge-2' }) + + await expect( + searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + workspaceId: 'workspace-1', + knowledgeBaseIds: ['knowledge-1', 'knowledge-2'], + topK: 5, + tagFilters: [{ tagName: 'team', operator: 'eq', value: 'docs' }], + }, + }) + ).rejects.toMatchObject({ code: 'validation' }) + + expect(mocks.generateEmbedding).not.toHaveBeenCalled() + expect(mocks.executeSearch).not.toHaveBeenCalled() + }) + + it('verifies trusted result provenance inside the authorized use case', async () => { + const registry = { markIncomplete: vi.fn() } + await searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + workspaceId: 'workspace-1', + knowledgeBaseIds: ['knowledge-1'], + query: 'answer', + topK: 5, + resultSecretRegistry: registry as never, + }, + }) + + expect(mocks.importProvenance).toHaveBeenCalledWith({ + registry, + results: expect.arrayContaining([ + expect.objectContaining({ id: 'embedding-1', documentId: 'document-1' }), + ]), + }) + }) + + it('propagates tag-definition infrastructure failures', async () => { + const failure = new Error('tag database unavailable') + mocks.getTagDefinitions.mockRejectedValueOnce(failure) + + await expect( + searchKnowledge.execute({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + input: { + workspaceId: 'workspace-1', + knowledgeBaseIds: ['knowledge-1'], + query: 'answer', + topK: 5, + }, + }) + ).rejects.toBe(failure) + }) +}) diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts new file mode 100644 index 00000000000..778806335df --- /dev/null +++ b/apps/sim/lib/knowledge/application/search.ts @@ -0,0 +1,278 @@ +import { checkAttributedUsageLimits } from '@/lib/billing/core/billing-attribution' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { generateRequestId } from '@/lib/core/utils/request' +import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' +import { + KnowledgeUsageLimitExceededError, + resolveKnowledgeAttributedUserId, + resolveKnowledgeBillingAttribution, +} from '@/lib/knowledge/application/billing' +import { + type KnowledgeWorkspaceContext, + resolveKnowledgeWorkspaceContext, +} from '@/lib/knowledge/application/contexts' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants' +import { recordSearchEmbeddingUsage } from '@/lib/knowledge/embeddings' +import { + executeKnowledgeSearch, + generateSearchEmbedding, + getDocumentMetadataByIds, + type SearchResult, +} from '@/lib/knowledge/search/queries' +import { importKnowledgeSearchResultSecretProvenance } from '@/lib/knowledge/secret-provenance' +import { getKnowledgeBaseById } from '@/lib/knowledge/service' +import { getDocumentTagDefinitions } from '@/lib/knowledge/tags/service' +import { buildUndefinedTagsError, validateTagValue } from '@/lib/knowledge/tags/utils' +import type { KnowledgeBaseWithCounts, StructuredFilter } from '@/lib/knowledge/types' +import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' + +export interface KnowledgeSearchTagFilter { + tagName: string + fieldType?: 'text' | 'number' | 'date' | 'boolean' + operator: string + value: string | number | boolean + valueTo?: string | number +} + +export interface SearchKnowledgeInput { + workspaceId: string + knowledgeBaseIds: string[] + query?: string + topK: number + tagFilters?: KnowledgeSearchTagFilter[] + /** Trusted execution provenance sink; never sourced from an HTTP or model payload. */ + resultSecretRegistry?: ResolvedSecretTraceRegistry +} + +interface KnowledgeSearchContext extends KnowledgeWorkspaceContext { + knowledgeBases: KnowledgeBaseWithCounts[] +} + +export interface KnowledgeSearchItem { + /** Trusted embedding identity for provenance import; HTTP presenters omit it. */ + embeddingId: string + documentId: string + documentName: string | null + sourceUrl: string | null + content: string + chunkIndex: number + metadata: Record + similarity: number +} + +export interface SearchKnowledgeResult { + results: KnowledgeSearchItem[] + query: string + knowledgeBaseIds: string[] + topK: number + totalResults: number +} + +async function resolveKnowledgeSearchContext( + input: SearchKnowledgeInput +): Promise { + if (input.knowledgeBaseIds.length < 1 || input.knowledgeBaseIds.length > 20) { + throw new OrchestrationError( + 'validation', + 'Knowledge search requires between 1 and 20 knowledge bases' + ) + } + if (!Number.isInteger(input.topK) || input.topK < 1 || input.topK > 100) { + throw new OrchestrationError('validation', 'topK must be an integer between 1 and 100') + } + const workspaceContext = await resolveKnowledgeWorkspaceContext(input) + const knowledgeBases = await Promise.all(input.knowledgeBaseIds.map(getKnowledgeBaseById)) + const inaccessibleIds = input.knowledgeBaseIds.filter( + (_id, index) => knowledgeBases[index]?.workspaceId !== workspaceContext.workspaceId + ) + if (inaccessibleIds.length > 0) { + throw new OrchestrationError( + 'not_found', + `Knowledge bases not found or access denied: ${inaccessibleIds.join(', ')}` + ) + } + return { + ...workspaceContext, + knowledgeBases: knowledgeBases as KnowledgeBaseWithCounts[], + } +} + +function buildStructuredFilters( + filters: KnowledgeSearchTagFilter[], + tagDefinitions: Awaited> +): StructuredFilter[] { + const definitionsByName = new Map( + tagDefinitions.map((definition) => [definition.displayName, definition]) + ) + const undefinedTags: string[] = [] + const typeErrors: string[] = [] + for (const filter of filters) { + const definition = definitionsByName.get(filter.tagName) + if (!definition) { + undefinedTags.push(filter.tagName) + continue + } + const validationError = validateTagValue( + filter.tagName, + String(filter.value), + definition.fieldType + ) + if (validationError) typeErrors.push(validationError) + } + if (undefinedTags.length > 0 || typeErrors.length > 0) { + const messages = [ + ...(undefinedTags.length > 0 ? [buildUndefinedTagsError(undefinedTags)] : []), + ...typeErrors, + ] + throw new OrchestrationError('validation', messages.join('\n')) + } + return filters.map((filter) => { + const definition = definitionsByName.get(filter.tagName) + if (!definition) throw new Error('Validated knowledge tag definition disappeared') + return { + tagSlot: definition.tagSlot, + fieldType: definition.fieldType, + operator: filter.operator, + value: filter.value, + valueTo: filter.valueTo, + } + }) +} + +export const searchKnowledge = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.search, + resolveContext: ({ input }: { input: SearchKnowledgeInput }) => + resolveKnowledgeSearchContext(input), + async execute({ principal, input, context }) { + const hasQuery = Boolean(input.query?.trim()) + const filters = input.tagFilters ?? [] + if (!hasQuery && filters.length === 0) { + throw new OrchestrationError('validation', 'Either query or tagFilters must be provided') + } + if (filters.length > 0 && context.knowledgeBases.length > 1) { + throw new OrchestrationError( + 'validation', + 'Tag filters are only supported when searching a single knowledge base' + ) + } + + const billingAttribution = hasQuery + ? await resolveKnowledgeBillingAttribution(principal, context) + : undefined + if (billingAttribution) { + const usage = await checkAttributedUsageLimits(billingAttribution) + if (usage.isExceeded) { + throw new KnowledgeUsageLimitExceededError( + usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' + ) + } + } + + const tagDefinitionsByKnowledgeBase = new Map< + string, + Awaited> + >() + let structuredFilters: StructuredFilter[] = [] + if (filters.length > 0) { + const knowledgeBaseId = context.knowledgeBases[0].id + const definitions = await getDocumentTagDefinitions(knowledgeBaseId) + tagDefinitionsByKnowledgeBase.set(knowledgeBaseId, definitions) + structuredFilters = buildStructuredFilters(filters, definitions) + } + + const embeddingModels = [...new Set(context.knowledgeBases.map((kb) => kb.embeddingModel))] + if (hasQuery && embeddingModels.length > 1) { + throw new OrchestrationError( + 'validation', + 'Selected knowledge bases use different embedding models and cannot be searched together. Search them separately.' + ) + } + const embeddingModel = embeddingModels[0] + let queryEmbeddingIsBYOK: boolean | null = null + let queryVector: string | undefined + if (hasQuery) { + const generated = await generateSearchEmbedding( + input.query!, + embeddingModel, + context.workspaceId + ) + queryEmbeddingIsBYOK = generated.isBYOK + queryVector = JSON.stringify(generated.embedding) + } + + const knowledgeBaseIds = context.knowledgeBases.map((kb) => kb.id) + const rows = await executeKnowledgeSearch({ + knowledgeBaseIds, + topK: input.topK, + searchMode: 'vector', + query: input.query, + queryVector, + structuredFilters, + }) + + if (input.resultSecretRegistry) { + const provenance = await importKnowledgeSearchResultSecretProvenance({ + registry: input.resultSecretRegistry, + results: rows, + }) + if (!provenance.imported) { + input.resultSecretRegistry.markIncomplete() + throw new Error('Knowledge result secret provenance is unavailable') + } + } + + if (queryEmbeddingIsBYOK !== null && billingAttribution) { + await recordSearchEmbeddingUsage({ + userId: resolveKnowledgeAttributedUserId(principal, context), + workspaceId: context.workspaceId, + embeddingModel, + query: input.query!, + isBYOK: queryEmbeddingIsBYOK, + sourceReference: `v2-kb-search:${generateRequestId()}`, + billingAttribution, + }) + } + + const tagDefinitionEntries = await Promise.all( + knowledgeBaseIds.map(async (knowledgeBaseId) => { + const definitions = + tagDefinitionsByKnowledgeBase.get(knowledgeBaseId) ?? + (await getDocumentTagDefinitions(knowledgeBaseId)) + return [ + knowledgeBaseId, + new Map(definitions.map((definition) => [definition.tagSlot, definition.displayName])), + ] as const + }) + ) + const tagMaps = new Map(tagDefinitionEntries) + const documentMetadata = await getDocumentMetadataByIds(rows.map((row) => row.documentId)) + + const results = rows.map((row: SearchResult): KnowledgeSearchItem => { + const metadata: Record = {} + const tagMap = tagMaps.get(row.knowledgeBaseId) + for (const slot of ALL_TAG_SLOTS) { + const value = row[slot] + if (value !== null && value !== undefined) metadata[tagMap?.get(slot) ?? slot] = value + } + const document = documentMetadata[row.documentId] + return { + embeddingId: row.id, + documentId: row.documentId, + documentName: document?.filename ?? null, + sourceUrl: document?.sourceUrl ?? null, + content: row.content, + chunkIndex: row.chunkIndex, + metadata, + similarity: hasQuery ? 1 - row.distance : 1, + } + }) + return { + results, + query: input.query ?? '', + knowledgeBaseIds, + topK: input.topK, + totalResults: results.length, + } + }, +}) diff --git a/apps/sim/lib/knowledge/application/upload-sessions.test.ts b/apps/sim/lib/knowledge/application/upload-sessions.test.ts new file mode 100644 index 00000000000..e4b43e9cf83 --- /dev/null +++ b/apps/sim/lib/knowledge/application/upload-sessions.test.ts @@ -0,0 +1,620 @@ +/** + * @vitest-environment node + */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + abortUpload: vi.fn(), + assertBinding: vi.fn(), + checkUsage: vi.fn(), + completeUpload: vi.fn(), + createDocument: vi.fn(), + createPartUrls: vi.fn(), + createUpload: vi.fn(), + findBound: vi.fn(), + getUpload: vi.fn(), + processQueue: vi.fn(), + recordAudit: vi.fn(), + recordOwnership: vi.fn(), + resolveBilling: vi.fn(), + resolveContext: vi.fn(), + resolvePermission: vi.fn(), + validateFileType: vi.fn(), +})) + +vi.mock('@sim/audit', () => ({ + AuditAction: { DOCUMENT_UPLOADED: 'document.uploaded' }, + AuditResourceType: { DOCUMENT: 'document' }, + recordAudit: mocks.recordAudit, +})) + +vi.mock('@sim/platform-authz/workspace', () => ({ + permissionSatisfies: (actual: string | null, required: string) => { + const rank = { read: 1, write: 2, admin: 3 } as const + return ( + actual !== null && rank[actual as keyof typeof rank] >= rank[required as keyof typeof rank] + ) + }, + resolveEffectiveWorkspacePermission: mocks.resolvePermission, +})) + +vi.mock('@/lib/billing/core/billing-attribution', () => ({ + checkAttributedUsageLimits: mocks.checkUsage, + resolveBillingAttribution: mocks.resolveBilling, + resolveSystemBillingAttribution: mocks.resolveBilling, +})) + +vi.mock('@/lib/knowledge/application/contexts', () => ({ + resolveActiveKnowledgeBaseContext: mocks.resolveContext, +})) + +vi.mock('@/lib/knowledge/documents/service', () => ({ + createSingleDocument: mocks.createDocument, + processDocumentsWithQueue: mocks.processQueue, +})) + +vi.mock('@/lib/knowledge/orchestration/documents', () => ({ + findBoundKnowledgeDocument: mocks.findBound, +})) + +vi.mock('@/lib/uploads/server/metadata', () => ({ + recordKnowledgeBaseFileOwnership: mocks.recordOwnership, +})) + +vi.mock('@/lib/uploads/upload-session/application', () => ({ + requestOrigin: () => 'http://localhost:3000', +})) + +vi.mock('@/lib/uploads/upload-session/service', () => ({ + abortUploadSession: mocks.abortUpload, + assertUploadSessionAuthBinding: mocks.assertBinding, + completeUploadSession: mocks.completeUpload, + createUploadPartUrls: mocks.createPartUrls, + createUploadSession: mocks.createUpload, + getPrincipalKnowledgeDocumentUploadSession: mocks.getUpload, +})) + +vi.mock('@/lib/uploads/utils/validation', () => ({ + validateFileType: mocks.validateFileType, +})) + +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { + cancelKnowledgeDocumentUpload, + completeKnowledgeDocumentUpload, + createKnowledgeDocumentUpload, + issueKnowledgeDocumentUploadParts, +} from '@/lib/knowledge/application/upload-sessions' +import type { UploadSessionRecord } from '@/lib/uploads/upload-session/service' + +const CONTEXT = { + workspaceId: 'workspace-1', + workspaceOrganizationId: 'organization-1', + allowPersonalApiKeys: true, + billedAccountUserId: 'billing-owner-1', + knowledgeBaseId: 'knowledge-1', + knowledgeBase: { id: 'knowledge-1', name: 'Docs', workspaceId: 'workspace-1' }, +} +const PRINCIPAL = { + kind: 'personal_api_key' as const, + userId: 'user-1', + keyId: 'key-1', +} +const BILLING = { + actorUserId: 'user-1', + workspaceId: 'workspace-1', + organizationId: 'organization-1', + billedAccountUserId: 'billing-owner-1', + billingEntity: { id: 'organization-1', type: 'organization' }, + billingPeriod: { start: '2026-08-01', end: '2026-09-01' }, + payerSubscription: null, +} +const SESSION: UploadSessionRecord = { + id: 'upload-1', + workspaceId: 'workspace-1', + userId: 'user-1', + knowledgeBaseId: 'knowledge-1', + workflowId: null, + executionId: null, + purpose: 'knowledge_document', + method: 'multipart', + storageContext: 'knowledge-base', + storageKey: 'kb/guide.pdf', + finalKey: 'kb/guide.pdf', + storageProvider: 's3', + providerUploadId: 'provider-1', + providerObjectVersion: null, + fileName: 'guide.pdf', + contentType: 'application/pdf', + fileSize: 1024, + partSize: 8 * 1024 * 1024, + partCount: 1, + status: 'uploading', + metadata: { + tag1: 'product', + processingOptions: { recipe: 'default', lang: 'en' }, + authBinding: { + version: 1, + workspaceId: 'workspace-1', + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + }, + }, + uploadToken: 'token', + createdAt: new Date('2026-08-03T21:00:00.000Z'), + expiresAt: new Date('2026-08-04T21:00:00.000Z'), + completedFileId: null, + error: null, + completedAt: null, + updatedAt: new Date('2026-08-03T21:00:00.000Z'), +} +const DOCUMENT = { + id: 'upload-1', + knowledgeBaseId: 'knowledge-1', + filename: 'guide.pdf', + fileUrl: '/api/files/serve/s3/kb%2Fguide.pdf?context=knowledge-base', + fileSize: 1024, + mimeType: 'application/pdf', + chunkCount: 0, + tokenCount: 0, + characterCount: 0, + enabled: true, + uploadedAt: new Date('2026-08-03T21:01:00.000Z'), + tag1: 'product', + tag2: null, + tag3: null, + tag4: null, + tag5: null, + tag6: null, + tag7: null, +} +const REQUEST = { headers: new Headers() } + +describe('knowledge-document upload application lifecycle', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.resolveContext.mockResolvedValue(CONTEXT) + mocks.resolvePermission.mockResolvedValue('write') + mocks.resolveBilling.mockResolvedValue(BILLING) + mocks.checkUsage.mockResolvedValue({ isExceeded: false }) + mocks.validateFileType.mockReturnValue(null) + mocks.createUpload.mockResolvedValue({ + ...SESSION, + transfer: { method: 'multipart', partSize: SESSION.partSize, partCount: 1 }, + }) + mocks.recordOwnership.mockResolvedValue(undefined) + mocks.getUpload.mockResolvedValue(SESSION) + mocks.createPartUrls.mockResolvedValue([ + { + partNumber: 1, + url: 'https://storage.example/1', + headers: {}, + expiresAt: '2026-08-04T21:00:00.000Z', + }, + ]) + mocks.abortUpload.mockResolvedValue({ ...SESSION, status: 'aborted' }) + mocks.findBound.mockResolvedValue({ status: 'absent' }) + mocks.createDocument.mockResolvedValue(DOCUMENT) + mocks.processQueue.mockResolvedValue(undefined) + }) + + it('admits, binds, and records ownership before returning upload credentials', async () => { + await createKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + metadata: { tag1: 'product' }, + }, + request: REQUEST, + }) + + expect(mocks.createUpload).toHaveBeenCalledWith( + expect.objectContaining({ + purpose: 'knowledge_document', + principal: PRINCIPAL, + userId: 'user-1', + workspaceId: 'workspace-1', + knowledgeBaseId: 'knowledge-1', + }) + ) + expect(mocks.recordOwnership).toHaveBeenCalledWith( + expect.objectContaining({ key: SESSION.storageKey, workspaceId: 'workspace-1' }) + ) + expect(mocks.recordOwnership.mock.invocationCallOrder[0]).toBeGreaterThan( + mocks.createUpload.mock.invocationCallOrder[0] + ) + }) + + it('rejects insufficient role before allocating provider state', async () => { + mocks.resolvePermission.mockResolvedValue('read') + + await expect( + createKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + metadata: {}, + }, + request: REQUEST, + }) + ).rejects.toMatchObject({ code: 'forbidden' }) + expect(mocks.createUpload).not.toHaveBeenCalled() + }) + + it('aborts provider state and propagates an ownership registration failure', async () => { + const failure = new Error('ownership database unavailable') + mocks.recordOwnership.mockRejectedValue(failure) + + await expect( + createKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + name: 'guide.pdf', + contentType: 'application/pdf', + size: 1024, + metadata: {}, + }, + request: REQUEST, + }) + ).rejects.toBe(failure) + expect(mocks.abortUpload).toHaveBeenCalledWith(expect.objectContaining({ id: 'upload-1' })) + }) + + it('reauthorizes and verifies the immutable credential on the parts leg', async () => { + await issueKnowledgeDocumentUploadParts.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + partNumbers: [1], + }, + request: REQUEST, + }) + + expect(mocks.getUpload).toHaveBeenCalledWith( + expect.objectContaining({ principal: PRINCIPAL, uploadId: 'upload-1' }) + ) + expect(mocks.assertBinding).toHaveBeenCalledWith(SESSION, PRINCIPAL) + expect(mocks.resolvePermission).toHaveBeenCalledTimes(2) + expect(mocks.createPartUrls).toHaveBeenCalledWith( + expect.objectContaining({ session: SESSION, partNumbers: [1] }) + ) + }) + + it('checks durable binding before canceling', async () => { + await cancelKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + }, + request: REQUEST, + }) + + expect(mocks.findBound).toHaveBeenCalledWith( + expect.objectContaining({ documentId: 'upload-1', knowledgeBaseId: 'knowledge-1' }) + ) + expect(mocks.abortUpload).toHaveBeenCalledWith(SESSION) + }) + + it('reauthorizes immediately before durable registration and audits the created document', async () => { + mocks.completeUpload.mockImplementation( + async (params: { + session: UploadSessionRecord + finalize: (session: UploadSessionRecord) => Promise<{ + value: { document: typeof DOCUMENT; created: boolean; knowledgeBaseName: string | null } + completedFileId?: string + }> + }) => { + const finalized = await params.finalize(params.session) + return { + session: { ...params.session, status: 'completed' as const }, + value: finalized.value, + alreadyCompleted: false, + } + } + ) + + const result = await completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + source: 'api', + }, + request: REQUEST, + }) + + expect(result.value.created).toBe(true) + expect(mocks.resolvePermission.mock.calls.length).toBeGreaterThanOrEqual(4) + expect(mocks.createDocument).toHaveBeenCalledWith( + expect.any(Object), + 'knowledge-1', + expect.any(String), + 'user-1', + 'upload-1', + undefined, + { expectedWorkspaceId: 'workspace-1' } + ) + expect(mocks.recordAudit).toHaveBeenCalledWith( + expect.objectContaining({ + action: 'document.uploaded', + resourceId: 'upload-1', + metadata: expect.objectContaining({ operation: 'knowledge.documents.upload.complete' }), + }) + ) + }) + + it('returns an already-bound document without re-billing, re-registering, or auditing', async () => { + mocks.findBound.mockResolvedValue({ status: 'bound', document: DOCUMENT }) + mocks.completeUpload.mockImplementation( + async (params: { + session: UploadSessionRecord + finalize: (session: UploadSessionRecord) => Promise<{ + value: { document: typeof DOCUMENT; created: boolean; knowledgeBaseName: string | null } + }> + }) => ({ + session: { ...params.session, status: 'completed' as const }, + value: (await params.finalize({ ...params.session, error: null })).value, + alreadyCompleted: true, + }) + ) + + const result = await completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + source: 'api', + }, + request: REQUEST, + }) + + expect(result.value.created).toBe(false) + expect(mocks.resolveBilling).not.toHaveBeenCalled() + expect(mocks.createDocument).not.toHaveBeenCalled() + expect(mocks.processQueue).not.toHaveBeenCalled() + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) + + it('fails completion when document processing cannot be dispatched', async () => { + const failure = new Error('queue unavailable') + mocks.processQueue.mockRejectedValue(failure) + mocks.completeUpload.mockImplementation( + async (params: { + session: UploadSessionRecord + finalize: (session: UploadSessionRecord) => Promise + }) => params.finalize(params.session) + ) + + await expect( + completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + source: 'api', + }, + request: REQUEST, + }) + ).rejects.toMatchObject({ + name: 'KnowledgeDocumentProcessingDispatchError', + message: 'Knowledge document processing dispatch failed', + cause: failure, + }) + expect(mocks.createDocument).toHaveBeenCalledTimes(1) + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) + + it('retries a failed processing dispatch before completing a bound registration', async () => { + const recoveringSession = { + ...SESSION, + status: 'finalizing' as const, + completedFileId: null, + error: 'Knowledge document processing dispatch failed', + } + mocks.getUpload.mockResolvedValue(recoveringSession) + mocks.findBound.mockResolvedValue({ + status: 'bound', + document: { ...DOCUMENT, processingStatus: 'pending' }, + }) + mocks.completeUpload.mockImplementation( + async (params: { + session: UploadSessionRecord + finalize: (session: UploadSessionRecord) => Promise<{ + value: { document: typeof DOCUMENT; created: boolean; knowledgeBaseName: string | null } + }> + }) => ({ + session: { ...params.session, status: 'completed' as const }, + value: (await params.finalize({ ...params.session, error: null })).value, + alreadyCompleted: true, + }) + ) + + const result = await completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + source: 'api', + }, + request: REQUEST, + }) + + expect(result.value.created).toBe(true) + expect(mocks.resolveBilling).toHaveBeenCalledTimes(1) + expect(mocks.processQueue).toHaveBeenCalledTimes(1) + expect(mocks.createDocument).not.toHaveBeenCalled() + expect(mocks.recordAudit).toHaveBeenCalledTimes(1) + }) + + it('converges a finalization retry after durable bind without duplicate document or audit', async () => { + const recoveringSession = { + ...SESSION, + status: 'finalizing' as const, + completedFileId: null, + } + const completedSession = { + ...recoveringSession, + status: 'completed' as const, + completedFileId: DOCUMENT.id, + } + mocks.getUpload.mockResolvedValueOnce(recoveringSession).mockResolvedValueOnce(completedSession) + mocks.findBound.mockResolvedValue({ status: 'bound', document: DOCUMENT }) + mocks.completeUpload.mockImplementation( + async (params: { + session: UploadSessionRecord + finalize: (session: UploadSessionRecord) => Promise<{ + value: { document: typeof DOCUMENT; created: boolean; knowledgeBaseName: string | null } + }> + loadCompleted: (session: UploadSessionRecord) => Promise<{ + document: typeof DOCUMENT + created: boolean + knowledgeBaseName: string | null + }> + }) => { + if (params.session.status === 'completed') { + return { + session: params.session, + value: await params.loadCompleted(params.session), + alreadyCompleted: true, + } + } + return { + session: completedSession, + value: (await params.finalize(params.session)).value, + alreadyCompleted: true, + } + } + ) + + const input = { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + source: 'api' as const, + } + const recovered = await completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input, + request: REQUEST, + }) + const retry = await completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input, + request: REQUEST, + }) + + expect(recovered.value.created).toBe(true) + expect(retry.value.created).toBe(false) + expect(mocks.createDocument).not.toHaveBeenCalled() + expect(mocks.processQueue).not.toHaveBeenCalled() + expect(mocks.recordAudit).toHaveBeenCalledTimes(1) + }) + + it('fails fast when billing ownership changes before durable registration', async () => { + mocks.resolveBilling.mockResolvedValue({ ...BILLING, billedAccountUserId: 'stale-owner' }) + mocks.completeUpload.mockImplementation( + async (params: { + session: UploadSessionRecord + finalize: (session: UploadSessionRecord) => Promise + }) => params.finalize(params.session) + ) + + await expect( + completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + source: 'api', + }, + request: REQUEST, + }) + ).rejects.toThrow('billing attribution changed') + expect(mocks.createDocument).not.toHaveBeenCalled() + }) + + it('propagates provider completion failures without audit or registration', async () => { + const failure = new Error('provider unavailable') + mocks.completeUpload.mockRejectedValue(failure) + + await expect( + completeKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + source: 'api', + }, + request: REQUEST, + }) + ).rejects.toBe(failure) + expect(mocks.createDocument).not.toHaveBeenCalled() + expect(mocks.recordAudit).not.toHaveBeenCalled() + }) + + it('propagates canonical infrastructure failures without concealment fallback', async () => { + const failure = new Error('database unavailable') + mocks.resolveContext.mockRejectedValue(failure) + + await expect( + issueKnowledgeDocumentUploadParts.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'workspace-1', + uploadId: 'upload-1', + uploadToken: 'token', + partNumbers: [1], + }, + request: REQUEST, + }) + ).rejects.toBe(failure) + }) + + it('conceals an asserted workspace mismatch as not found', async () => { + mocks.resolveContext.mockRejectedValue( + new OrchestrationError('not_found', 'Knowledge base not found') + ) + + await expect( + cancelKnowledgeDocumentUpload.execute({ + principal: PRINCIPAL, + input: { + knowledgeBaseId: 'knowledge-1', + assertedWorkspaceId: 'different-workspace', + uploadId: 'upload-1', + uploadToken: 'token', + }, + request: REQUEST, + }) + ).rejects.toMatchObject({ code: 'not_found' }) + }) +}) diff --git a/apps/sim/lib/knowledge/application/upload-sessions.ts b/apps/sim/lib/knowledge/application/upload-sessions.ts new file mode 100644 index 00000000000..3d8586ad9da --- /dev/null +++ b/apps/sim/lib/knowledge/application/upload-sessions.ts @@ -0,0 +1,471 @@ +import { AuditAction, AuditResourceType } from '@sim/audit' +import type { Principal } from '@sim/auth/principal' +import type { V2KnowledgeDocumentUploadMetadata } from '@/lib/api/contracts/v2/knowledge' +import { v2KnowledgeDocumentUploadMetadataSchema } from '@/lib/api/contracts/v2/knowledge' +import { checkAttributedUsageLimits } from '@/lib/billing/core/billing-attribution' +import { authorizeWorkspaceOperation, type WorkspaceOperation } from '@/lib/core/application' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { generateRequestId } from '@/lib/core/utils/request' +import { knowledgeDelegationPolicy } from '@/lib/knowledge/application/authorization' +import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' +import { + KnowledgeUsageLimitExceededError, + resolveKnowledgeAttributedUserId, + resolveKnowledgeBillingAttribution, +} from '@/lib/knowledge/application/billing' +import { + type ActiveKnowledgeBaseContext, + resolveActiveKnowledgeBaseContext, +} from '@/lib/knowledge/application/contexts' +import { knowledgeOperations } from '@/lib/knowledge/application/operations' +import { + createSingleDocument, + type DocumentData, + processDocumentsWithQueue, +} from '@/lib/knowledge/documents/service' +import type { CreatedKnowledgeDocument } from '@/lib/knowledge/orchestration/documents' +import { findBoundKnowledgeDocument } from '@/lib/knowledge/orchestration/documents' +import { recordKnowledgeBaseFileOwnership } from '@/lib/uploads/server/metadata' +import { requestOrigin } from '@/lib/uploads/upload-session/application' +import { + abortUploadSession, + assertUploadSessionAuthBinding, + completeUploadSession, + createUploadPartUrls, + createUploadSession, + getPrincipalKnowledgeDocumentUploadSession, + type UploadSessionRecord, +} from '@/lib/uploads/upload-session/service' +import { validateFileType } from '@/lib/uploads/utils/validation' + +const PROCESSING_DISPATCH_FAILURE_MESSAGE = 'Knowledge document processing dispatch failed' + +class KnowledgeDocumentProcessingDispatchError extends Error { + constructor(cause: unknown) { + super(PROCESSING_DISPATCH_FAILURE_MESSAGE, { cause }) + this.name = 'KnowledgeDocumentProcessingDispatchError' + } +} + +export class KnowledgeDocumentUnsupportedMediaTypeError extends Error { + constructor(message: string) { + super(message) + this.name = 'KnowledgeDocumentUnsupportedMediaTypeError' + } +} + +export interface CreateKnowledgeDocumentUploadInput { + knowledgeBaseId: string + assertedWorkspaceId: string + name: string + contentType: string + size: number + metadata: V2KnowledgeDocumentUploadMetadata +} + +export interface KnowledgeDocumentUploadControlInput { + knowledgeBaseId: string + assertedWorkspaceId: string + uploadId: string + uploadToken: string +} + +export interface IssueKnowledgeDocumentUploadPartsInput + extends KnowledgeDocumentUploadControlInput { + partNumbers: number[] +} + +export interface CompleteKnowledgeDocumentUploadInput extends KnowledgeDocumentUploadControlInput { + source: 'api' | 'ui' +} + +interface KnowledgeDocumentUploadCompletion { + document: CreatedKnowledgeDocument + created: boolean + knowledgeBaseName: string | null +} + +export interface CompleteKnowledgeDocumentUploadResult { + session: UploadSessionRecord + value: KnowledgeDocumentUploadCompletion + alreadyCompleted: boolean + workspaceId: string + knowledgeBaseId: string +} + +export const createKnowledgeDocumentUpload = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.uploadCreate, + resolveContext: ({ input }: { input: CreateKnowledgeDocumentUploadInput }) => + resolveActiveKnowledgeBaseContext({ + knowledgeBaseId: input.knowledgeBaseId, + assertedWorkspaceId: input.assertedWorkspaceId, + }), + async execute({ principal, input, context, request }) { + if (!request) throw new Error('Knowledge upload creation requires a request context') + const billingAttribution = await resolveKnowledgeBillingAttribution(principal, context) + const usage = await checkAttributedUsageLimits(billingAttribution) + if (usage.isExceeded) { + throw new KnowledgeUsageLimitExceededError( + usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' + ) + } + + const fileTypeError = validateFileType(input.name, input.contentType) + if (fileTypeError) { + throw new KnowledgeDocumentUnsupportedMediaTypeError(fileTypeError.message) + } + + const storageActorUserId = resolveKnowledgeAttributedUserId(principal, context) + const session = await createUploadSession({ + purpose: 'knowledge_document', + workspaceId: context.workspaceId, + knowledgeBaseId: context.knowledgeBaseId, + userId: storageActorUserId, + principal, + fileName: input.name, + contentType: input.contentType, + fileSize: input.size, + metadata: input.metadata, + localOrigin: requestOrigin(request), + }) + try { + await recordKnowledgeBaseFileOwnership({ + key: session.storageKey, + userId: storageActorUserId, + workspaceId: context.workspaceId, + originalName: input.name, + contentType: input.contentType, + size: input.size, + }) + } catch (error) { + await abortUploadSession(session) + throw error + } + return session + }, +}) + +export const issueKnowledgeDocumentUploadParts = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.uploadParts, + resolveContext: ({ input }: { input: IssueKnowledgeDocumentUploadPartsInput }) => + resolveActiveKnowledgeBaseContext({ + knowledgeBaseId: input.knowledgeBaseId, + assertedWorkspaceId: input.assertedWorkspaceId, + }), + async execute({ principal, input, context, request }) { + if (!request) throw new Error('Knowledge upload part issuance requires a request context') + const session = await loadBoundKnowledgeDocumentUpload(principal, input, context) + await reauthorizeKnowledgeDocumentUpload(principal, session, knowledgeOperations.uploadParts) + return { + parts: await createUploadPartUrls({ + session, + partNumbers: input.partNumbers, + localOrigin: requestOrigin(request), + }), + } + }, +}) + +export const cancelKnowledgeDocumentUpload = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.uploadCancel, + resolveContext: ({ input }: { input: KnowledgeDocumentUploadControlInput }) => + resolveActiveKnowledgeBaseContext({ + knowledgeBaseId: input.knowledgeBaseId, + assertedWorkspaceId: input.assertedWorkspaceId, + }), + async execute({ principal, input, context }) { + const session = await loadBoundKnowledgeDocumentUpload(principal, input, context) + await reauthorizeKnowledgeDocumentUpload(principal, session, knowledgeOperations.uploadCancel) + const bound = await findBoundKnowledgeDocument({ + documentId: session.id, + knowledgeBaseId: context.knowledgeBaseId, + document: knowledgeDocumentInputFor(session), + }) + if (bound.status !== 'absent') { + throw new OrchestrationError('conflict', 'Upload has already been completed') + } + return abortUploadSession(session) + }, +}) + +export const completeKnowledgeDocumentUpload = defineAuthorizedKnowledgeUseCase({ + operation: knowledgeOperations.uploadComplete, + resolveContext: ({ input }: { input: CompleteKnowledgeDocumentUploadInput }) => + resolveActiveKnowledgeBaseContext({ + knowledgeBaseId: input.knowledgeBaseId, + assertedWorkspaceId: input.assertedWorkspaceId, + }), + async execute({ + principal, + input, + context, + request, + }): Promise { + if (!request) throw new Error('Knowledge upload completion requires a request context') + const session = await loadBoundKnowledgeDocumentUpload(principal, input, context) + await reauthorizeKnowledgeDocumentUpload(principal, session, knowledgeOperations.uploadComplete) + const requestId = generateRequestId() + const recoveringUnprojectedRegistration = + session.status === 'finalizing' && session.completedFileId === null + const result = await completeUploadSession({ + session, + loadCompleted: async (claimed) => { + const freshContext = await reauthorizeKnowledgeDocumentUpload( + principal, + claimed, + knowledgeOperations.uploadComplete + ) + const document = knowledgeDocumentInputFor(claimed) + const bound = await findBoundKnowledgeDocument({ + documentId: claimed.id, + knowledgeBaseId: freshContext.knowledgeBaseId, + document, + }) + if (bound.status === 'conflict') { + throw new OrchestrationError( + 'conflict', + 'Upload id is already bound to a different document' + ) + } + if (bound.status === 'absent') { + throw new Error('Completed knowledge upload is missing its durable document') + } + if (claimed.completedFileId !== bound.document.id) { + throw new Error('Completed knowledge upload references a different durable document') + } + return { + document: bound.document, + created: false, + knowledgeBaseName: freshContext.knowledgeBase.name, + } + }, + finalize: async (claimed) => { + const freshContext = await reauthorizeKnowledgeDocumentUpload( + principal, + claimed, + knowledgeOperations.uploadComplete + ) + const { processingOptions } = knowledgeDocumentMetadataFor(claimed) + const document = knowledgeDocumentInputFor(claimed) + const bound = await findBoundKnowledgeDocument({ + documentId: claimed.id, + knowledgeBaseId: freshContext.knowledgeBaseId, + document, + }) + if (bound.status === 'conflict') { + throw new OrchestrationError( + 'conflict', + 'Upload id is already bound to a different document' + ) + } + if (bound.status === 'bound') { + if ( + session.error === PROCESSING_DISPATCH_FAILURE_MESSAGE && + bound.document.processingStatus === 'pending' + ) { + const billingAttribution = await resolveKnowledgeBillingAttribution( + principal, + freshContext + ) + await dispatchKnowledgeDocumentProcessing( + bound.document, + freshContext.knowledgeBaseId, + processingOptions, + requestId, + billingAttribution + ) + } + return { + value: { + document: bound.document, + created: recoveringUnprojectedRegistration, + knowledgeBaseName: freshContext.knowledgeBase.name, + }, + completedFileId: bound.document.id, + } + } + + const billingAttribution = await resolveKnowledgeBillingAttribution(principal, freshContext) + const registrationContext = await reauthorizeKnowledgeDocumentUpload( + principal, + claimed, + knowledgeOperations.uploadComplete + ) + const uploadedBy = resolveKnowledgeAttributedUserId(principal, registrationContext) + if ( + billingAttribution.workspaceId !== registrationContext.workspaceId || + billingAttribution.actorUserId !== uploadedBy || + billingAttribution.organizationId !== registrationContext.workspaceOrganizationId || + billingAttribution.billedAccountUserId !== registrationContext.billedAccountUserId + ) { + throw new Error('Knowledge upload billing attribution changed before registration') + } + + let created: CreatedKnowledgeDocument + try { + created = await createSingleDocument( + document, + registrationContext.knowledgeBaseId, + requestId, + uploadedBy, + claimed.id, + undefined, + { expectedWorkspaceId: registrationContext.workspaceId } + ) + } catch (error) { + const afterError = await findBoundKnowledgeDocument({ + documentId: claimed.id, + knowledgeBaseId: registrationContext.knowledgeBaseId, + document, + }) + if (afterError.status === 'conflict') { + throw new OrchestrationError( + 'conflict', + 'Upload id is already bound to a different document' + ) + } + if (afterError.status === 'bound') { + return { + value: { + document: afterError.document, + created: false, + knowledgeBaseName: registrationContext.knowledgeBase.name, + }, + completedFileId: afterError.document.id, + } + } + throw error + } + + await dispatchKnowledgeDocumentProcessing( + created, + registrationContext.knowledgeBaseId, + processingOptions, + requestId, + billingAttribution + ) + return { + value: { + document: created, + created: true, + knowledgeBaseName: registrationContext.knowledgeBase.name, + }, + completedFileId: created.id, + } + }, + }) + return { + ...result, + workspaceId: context.workspaceId, + knowledgeBaseId: context.knowledgeBaseId, + } + }, + projectAudit: ({ input, context, result }) => { + if (!result.value.created) return [] + const { document, knowledgeBaseName } = result.value + return { + action: AuditAction.DOCUMENT_UPLOADED, + resourceType: AuditResourceType.DOCUMENT, + resourceId: document.id, + resourceName: document.filename, + description: `Uploaded document "${document.filename}" to knowledge base "${knowledgeBaseName ?? context.knowledgeBaseId}"`, + metadata: { + source: input.source, + knowledgeBaseId: context.knowledgeBaseId, + knowledgeBaseName, + fileName: document.filename, + fileType: document.mimeType, + fileSize: document.fileSize, + }, + } + }, +}) + +async function dispatchKnowledgeDocumentProcessing( + document: CreatedKnowledgeDocument, + knowledgeBaseId: string, + processingOptions: V2KnowledgeDocumentUploadMetadata['processingOptions'], + requestId: string, + billingAttribution: Awaited> +): Promise { + const processingDocument: DocumentData = { + documentId: document.id, + filename: document.filename, + fileUrl: document.fileUrl, + fileSize: document.fileSize, + mimeType: document.mimeType, + } + try { + await processDocumentsWithQueue( + [processingDocument], + knowledgeBaseId, + processingOptions ?? {}, + requestId, + billingAttribution + ) + } catch (error) { + throw new KnowledgeDocumentProcessingDispatchError(error) + } +} + +async function loadBoundKnowledgeDocumentUpload( + principal: Principal, + input: KnowledgeDocumentUploadControlInput, + context: ActiveKnowledgeBaseContext +): Promise { + return getPrincipalKnowledgeDocumentUploadSession({ + uploadId: input.uploadId, + uploadToken: input.uploadToken, + principal, + workspaceId: context.workspaceId, + knowledgeBaseId: context.knowledgeBaseId, + }) +} + +async function reauthorizeKnowledgeDocumentUpload( + principal: Principal, + session: UploadSessionRecord, + operation: WorkspaceOperation +): Promise { + if ( + !session.workspaceId || + !session.knowledgeBaseId || + session.purpose !== 'knowledge_document' + ) { + throw new OrchestrationError('not_found', 'Upload session not found') + } + assertUploadSessionAuthBinding(session, principal) + const context = await resolveActiveKnowledgeBaseContext({ + knowledgeBaseId: session.knowledgeBaseId, + assertedWorkspaceId: session.workspaceId, + }) + await authorizeWorkspaceOperation(principal, operation, context, { + delegation: knowledgeDelegationPolicy, + }) + return context +} + +function knowledgeDocumentMetadataFor(session: UploadSessionRecord) { + const { authBinding: _authBinding, ...metadata } = session.metadata + return v2KnowledgeDocumentUploadMetadataSchema.parse(metadata) +} + +function knowledgeDocumentInputFor(session: UploadSessionRecord) { + const { processingOptions: _processingOptions, ...documentTags } = + knowledgeDocumentMetadataFor(session) + return { + filename: session.fileName, + fileUrl: knowledgeDocumentFileUrl(session), + fileSize: session.fileSize, + mimeType: session.contentType, + ...documentTags, + } +} + +export function knowledgeDocumentFileUrl(session: UploadSessionRecord): string { + if (session.storageContext !== 'knowledge-base') { + throw new Error('Knowledge-document upload has an invalid storage context') + } + const providerPrefix = session.storageProvider === 'local' ? '' : `${session.storageProvider}/` + return `/api/files/serve/${providerPrefix}${encodeURIComponent(session.storageKey)}?context=knowledge-base` +} diff --git a/apps/sim/lib/knowledge/constants.ts b/apps/sim/lib/knowledge/constants.ts index 57ee50321be..e0f53db00ae 100644 --- a/apps/sim/lib/knowledge/constants.ts +++ b/apps/sim/lib/knowledge/constants.ts @@ -1,5 +1,11 @@ /** Max character length for a knowledge base description, enforced at every layer (UI, internal API, v1 API). */ export const KNOWLEDGE_BASE_DESCRIPTION_MAX_LENGTH = 10_000 +/** Hard bound for full-workspace knowledge-base list projections. */ +export const MAX_KNOWLEDGE_BASES_PER_WORKSPACE = 10_000 +/** Hard bound for path-indexed knowledge folder trees and recursive cascades. */ +export const MAX_KNOWLEDGE_FOLDERS_PER_WORKSPACE = 10_000 +/** Hard bound for connector-type rows projected onto one knowledge-base list. */ +export const MAX_KNOWLEDGE_CONNECTOR_TYPE_ROWS_PER_LIST = 100_000 /** * Chunking a knowledge base gets when its creator names no configuration. diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 187ef80076c..7425819f571 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -1838,7 +1838,8 @@ export async function createSingleDocument( requestId: string, uploadedBy: string | null = null, documentId = generateId(), - secretProvenance?: KnowledgeDocumentWriteSecretProvenance + secretProvenance?: KnowledgeDocumentWriteSecretProvenance, + options?: { expectedWorkspaceId?: string } ): Promise<{ id: string knowledgeBaseId: string @@ -1939,6 +1940,13 @@ export async function createSingleDocument( throw new OrchestrationError('not_found', 'Knowledge base not found') } + if ( + options?.expectedWorkspaceId !== undefined && + kb[0].workspaceId !== options.expectedWorkspaceId + ) { + throw new OrchestrationError('not_found', 'Knowledge base not found') + } + if ( kb[0].workspaceId !== admission.workspaceId || kb[0].userId !== admission.knowledgeBaseUserId @@ -2746,7 +2754,8 @@ async function excludeConnectorDocuments( async function deleteDocumentsByLifecyclePolicy( documentIds: string[], - requestId: string + requestId: string, + expectedKnowledgeBaseId?: string ): Promise { const ids = [...new Set(documentIds)] if (ids.length === 0) { @@ -2759,14 +2768,26 @@ async function deleteDocumentsByLifecyclePolicy( connectorId: document.connectorId, }) .from(document) - .where(inArray(document.id, ids)) + .where( + expectedKnowledgeBaseId + ? and( + inArray(document.id, ids), + eq(document.knowledgeBaseId, expectedKnowledgeBaseId), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) + : inArray(document.id, ids) + ) const connectorBackedIds = docs.filter((doc) => doc.connectorId !== null).map((doc) => doc.id) const hardDeleteIds = docs.filter((doc) => doc.connectorId === null).map((doc) => doc.id) const [excludedCount, hardDeletedCount] = await Promise.all([ - excludeConnectorDocuments(connectorBackedIds, requestId), - hardDeleteDocuments(hardDeleteIds, requestId), + expectedKnowledgeBaseId + ? excludeConnectorKnowledgeDocuments(expectedKnowledgeBaseId, connectorBackedIds, requestId) + : excludeConnectorDocuments(connectorBackedIds, requestId), + hardDeleteDocuments(hardDeleteIds, requestId, undefined, expectedKnowledgeBaseId), ]) return excludedCount + hardDeletedCount @@ -2784,7 +2805,8 @@ export async function hardDeleteDocuments( * connector, keep documents") would otherwise still have them purged here * despite no longer belonging to the connector the caller reasoned about. */ - expectedConnectorId?: string + expectedConnectorId?: string, + expectedKnowledgeBaseId?: string ): Promise { const ids = [...new Set(documentIds)] if (ids.length === 0) { @@ -2796,7 +2818,8 @@ export async function hardDeleteDocuments( deletedCount += await hardDeleteDocumentBatch( ids.slice(offset, offset + HARD_DELETE_DOCUMENT_BATCH_SIZE), requestId, - expectedConnectorId + expectedConnectorId, + expectedKnowledgeBaseId ) } return deletedCount @@ -2809,7 +2832,8 @@ export async function hardDeleteDocuments( async function hardDeleteDocumentBatch( documentIds: string[], requestId: string, - expectedConnectorId?: string + expectedConnectorId?: string, + expectedKnowledgeBaseId?: string ): Promise { const ids = [...new Set(documentIds)] const documentsToDelete = await db @@ -2826,9 +2850,14 @@ async function hardDeleteDocumentBatch( .from(document) .innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id)) .where( - expectedConnectorId - ? and(inArray(document.id, ids), eq(document.connectorId, expectedConnectorId)) - : inArray(document.id, ids) + and( + inArray(document.id, ids), + expectedConnectorId ? eq(document.connectorId, expectedConnectorId) : undefined, + expectedKnowledgeBaseId ? eq(document.knowledgeBaseId, expectedKnowledgeBaseId) : undefined, + expectedKnowledgeBaseId ? eq(document.userExcluded, false) : undefined, + expectedKnowledgeBaseId ? isNull(document.archivedAt) : undefined, + expectedKnowledgeBaseId ? isNull(document.deletedAt) : undefined + ) ) if (documentsToDelete.length === 0) { @@ -2910,16 +2939,26 @@ async function hardDeleteDocumentBatch( * embedding delete and the document delete are scoped to this re-verified * ID set rather than the stale `existingIds`. */ - const stillTargetedIds = expectedConnectorId - ? ( - await tx - .select({ id: document.id }) - .from(document) - .where( - and(inArray(document.id, existingIds), eq(document.connectorId, expectedConnectorId)) - ) - ).map((d) => d.id) - : existingIds + const stillTargetedIds = + expectedConnectorId || expectedKnowledgeBaseId + ? ( + await tx + .select({ id: document.id }) + .from(document) + .where( + and( + inArray(document.id, existingIds), + expectedConnectorId ? eq(document.connectorId, expectedConnectorId) : undefined, + expectedKnowledgeBaseId + ? eq(document.knowledgeBaseId, expectedKnowledgeBaseId) + : undefined, + expectedKnowledgeBaseId ? eq(document.userExcluded, false) : undefined, + expectedKnowledgeBaseId ? isNull(document.archivedAt) : undefined, + expectedKnowledgeBaseId ? isNull(document.deletedAt) : undefined + ) + ) + ).map((d) => d.id) + : existingIds await tx.delete(embedding).where(inArray(embedding.documentId, stillTargetedIds)) const deletedRows = await tx @@ -2985,3 +3024,44 @@ export async function deleteDocument( message: 'Document deleted successfully', } } + +/** Deletes one currently visible document within its canonical knowledge base. */ +export async function deleteKnowledgeDocumentInKnowledgeBase( + knowledgeBaseId: string, + documentId: string, + requestId: string +): Promise { + const current = await getKnowledgeDocument(knowledgeBaseId, documentId) + if (!current) throw new OrchestrationError('not_found', 'Document not found') + const affected = await deleteDocumentsByLifecyclePolicy([documentId], requestId, knowledgeBaseId) + if (affected !== 1) throw new OrchestrationError('not_found', 'Document not found') +} + +async function excludeConnectorKnowledgeDocuments( + knowledgeBaseId: string, + documentIds: string[], + requestId: string +): Promise { + if (documentIds.length === 0) return 0 + const updated = await db + .update(document) + .set({ userExcluded: true, enabled: false }) + .where( + and( + inArray(document.id, documentIds), + eq(document.knowledgeBaseId, knowledgeBaseId), + isNotNull(document.connectorId), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) + ) + .returning({ id: document.id }) + if (updated.length > 0) { + logger.info(`[${requestId}] Excluded ${updated.length} connector-backed document(s)`, { + documentIds: updated.map((row) => row.id), + knowledgeBaseId, + }) + } + return updated.length +} diff --git a/apps/sim/lib/knowledge/service.test.ts b/apps/sim/lib/knowledge/service.test.ts index 7399b83aebe..b1be80ee910 100644 --- a/apps/sim/lib/knowledge/service.test.ts +++ b/apps/sim/lib/knowledge/service.test.ts @@ -31,7 +31,32 @@ vi.mock('@/lib/billing/core/usage', () => ({ ensureUserStatsExists: mockEnsureUserStatsExists, })) -import { KnowledgeBasePermissionError, updateKnowledgeBase } from '@/lib/knowledge/service' +import { MAX_KNOWLEDGE_BASES_PER_WORKSPACE } from '@/lib/knowledge/constants' +import { + getWorkspaceKnowledgeBases, + KnowledgeBasePermissionError, + updateKnowledgeBase, +} from '@/lib/knowledge/service' + +describe('getWorkspaceKnowledgeBases — bounded reads', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + it('fails before projecting connector data for an oversized workspace list', async () => { + dbChainMockFns.limit.mockResolvedValueOnce( + Array.from({ length: MAX_KNOWLEDGE_BASES_PER_WORKSPACE + 1 }, (_, index) => ({ + id: `kb-${index}`, + })) + ) + + await expect(getWorkspaceKnowledgeBases('ws-1')).rejects.toThrow( + `Knowledge base list exceeds the ${MAX_KNOWLEDGE_BASES_PER_WORKSPACE} row limit` + ) + expect(dbChainMockFns.limit).toHaveBeenCalledWith(MAX_KNOWLEDGE_BASES_PER_WORKSPACE + 1) + }) +}) /** * These tests guard the workspace mass-assignment fix: diff --git a/apps/sim/lib/knowledge/service.ts b/apps/sim/lib/knowledge/service.ts index 0991e17a5f3..8de2d573357 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -38,6 +38,10 @@ import { import { OrchestrationError } from '@/lib/core/orchestration/types' import { generateRestoreName } from '@/lib/core/utils/restore-name' import { findActiveFolder, resolveRestoredFolderId } from '@/lib/folders/queries' +import { + MAX_KNOWLEDGE_BASES_PER_WORKSPACE, + MAX_KNOWLEDGE_CONNECTOR_TYPE_ROWS_PER_LIST, +} from '@/lib/knowledge/constants' import type { ChunkingConfig, CreateKnowledgeBaseData, @@ -133,9 +137,8 @@ const KNOWLEDGE_BASE_SORTS = { updatedAt: [knowledgeBase.updatedAt, knowledgeBase.createdAt], } satisfies Record -interface GetKnowledgeBasesOptions { - /** Restrict to one knowledge-base folder. */ - /** `undefined` lists every folder, `null` lists only workspace-root resources. */ +export interface GetKnowledgeBasesOptions { + /** Restrict to one knowledge-base folder; `undefined` lists all and `null` lists the root. */ folderId?: string | null /** Case-insensitive substring match on the knowledge base name. */ search?: string @@ -143,6 +146,122 @@ interface GetKnowledgeBasesOptions { sortOrder?: V2SortOrder } +async function attachConnectorTypes( + knowledgeBases: Array> +): Promise { + const kbIds = knowledgeBases.map((kb) => kb.id) + const connectorRows = + kbIds.length > 0 + ? await db + .select({ + knowledgeBaseId: knowledgeConnector.knowledgeBaseId, + connectorType: knowledgeConnector.connectorType, + }) + .from(knowledgeConnector) + .where( + and( + inArray(knowledgeConnector.knowledgeBaseId, kbIds), + isNull(knowledgeConnector.archivedAt), + isNull(knowledgeConnector.deletedAt) + ) + ) + .limit(MAX_KNOWLEDGE_CONNECTOR_TYPE_ROWS_PER_LIST + 1) + : [] + if (connectorRows.length > MAX_KNOWLEDGE_CONNECTOR_TYPE_ROWS_PER_LIST) { + throw new Error( + `Knowledge connector projection exceeds the ${MAX_KNOWLEDGE_CONNECTOR_TYPE_ROWS_PER_LIST} row limit` + ) + } + + const connectorTypesByKb = new Map() + for (const row of connectorRows) { + const types = connectorTypesByKb.get(row.knowledgeBaseId) ?? [] + if (!types.includes(row.connectorType)) types.push(row.connectorType) + connectorTypesByKb.set(row.knowledgeBaseId, types) + } + + return knowledgeBases.map((kb) => ({ + ...kb, + connectorTypes: connectorTypesByKb.get(kb.id) ?? [], + })) +} + +/** + * Lists active knowledge bases in one canonical workspace after application + * authorization. Unlike the legacy user-oriented query, this never widens the + * scope to workspace-less rows and never depends on a human permission join. + */ +export async function getWorkspaceKnowledgeBases( + workspaceId: string, + scope: KnowledgeBaseScope = 'active', + options?: GetKnowledgeBasesOptions +): Promise { + const { folderId, search, sortBy = 'createdAt', sortOrder = 'asc' } = options ?? {} + const scopeCondition = + scope === 'all' + ? undefined + : scope === 'archived' + ? sql`${knowledgeBase.deletedAt} IS NOT NULL` + : isNull(knowledgeBase.deletedAt) + + const rows = await db + .select({ + id: knowledgeBase.id, + userId: knowledgeBase.userId, + name: knowledgeBase.name, + description: knowledgeBase.description, + tokenCount: sql`COALESCE(SUM(${document.tokenCount}), 0)`.mapWith(Number), + embeddingModel: knowledgeBase.embeddingModel, + embeddingDimension: knowledgeBase.embeddingDimension, + chunkingConfig: knowledgeBase.chunkingConfig, + createdAt: knowledgeBase.createdAt, + updatedAt: knowledgeBase.updatedAt, + deletedAt: knowledgeBase.deletedAt, + workspaceId: knowledgeBase.workspaceId, + folderId: knowledgeBase.folderId, + docCount: count(document.id), + }) + .from(knowledgeBase) + .leftJoin( + document, + and( + eq(document.knowledgeBaseId, knowledgeBase.id), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) + ) + .where( + and( + eq(knowledgeBase.workspaceId, workspaceId), + scopeCondition, + folderId === undefined + ? undefined + : folderId === null + ? isNull(knowledgeBase.folderId) + : eq(knowledgeBase.folderId, folderId), + searchFilter(knowledgeBase.name, search) + ) + ) + .groupBy(knowledgeBase.id) + .orderBy(...listOrderBy(KNOWLEDGE_BASE_SORTS[sortBy], sortOrder)) + .limit(MAX_KNOWLEDGE_BASES_PER_WORKSPACE + 1) + + if (rows.length > MAX_KNOWLEDGE_BASES_PER_WORKSPACE) { + throw new Error( + `Knowledge base list exceeds the ${MAX_KNOWLEDGE_BASES_PER_WORKSPACE} row limit` + ) + } + + return attachConnectorTypes( + rows.map((kb) => ({ + ...kb, + chunkingConfig: kb.chunkingConfig as ChunkingConfig, + docCount: Number(kb.docCount), + })) + ) +} + /** * Get knowledge bases that a user can access. * @@ -275,9 +394,6 @@ export async function createKnowledgeBase( data: CreateKnowledgeBaseData, requestId: string ): Promise { - const kbId = generateId() - const now = new Date() - const hasPermission = await getUserEntityPermissions(data.userId, 'workspace', data.workspaceId) if (hasPermission !== 'admin' && hasPermission !== 'write') { throw new KnowledgeBasePermissionError( @@ -285,6 +401,20 @@ export async function createKnowledgeBase( ) } + return createAuthorizedKnowledgeBase(data, requestId) +} + +/** + * Persists a knowledge base for an already-authorized application use case. + * Callers outside the application layer must use {@link createKnowledgeBase}. + */ +export async function createAuthorizedKnowledgeBase( + data: CreateKnowledgeBaseData, + requestId: string +): Promise { + const kbId = generateId() + const now = new Date() + await assertKnowledgeBaseFolder(data.folderId, data.workspaceId) const folderId = data.folderId ?? null @@ -368,7 +498,7 @@ export async function updateKnowledgeBase( } }, requestId: string, - options?: { actorUserId?: string } + options?: { actorUserId?: string; assertedWorkspaceId?: string } ): Promise { const now = new Date() const updateData: Partial = { @@ -403,7 +533,15 @@ export async function updateKnowledgeBase( const [snapshot] = await db .select({ workspaceId: knowledgeBase.workspaceId }) .from(knowledgeBase) - .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt))) + .where( + and( + eq(knowledgeBase.id, knowledgeBaseId), + isNull(knowledgeBase.deletedAt), + options?.assertedWorkspaceId + ? eq(knowledgeBase.workspaceId, options.assertedWorkspaceId) + : undefined + ) + ) .limit(1) if (!snapshot) { throw new KnowledgeBaseNotFoundError(knowledgeBaseId) @@ -427,7 +565,15 @@ export async function updateKnowledgeBase( folderId: knowledgeBase.folderId, }) .from(knowledgeBase) - .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt))) + .where( + and( + eq(knowledgeBase.id, knowledgeBaseId), + isNull(knowledgeBase.deletedAt), + options?.assertedWorkspaceId + ? eq(knowledgeBase.workspaceId, options.assertedWorkspaceId) + : undefined + ) + ) .limit(1) if (!kbSnapshot) { throw new KnowledgeBaseNotFoundError(knowledgeBaseId) @@ -510,7 +656,15 @@ export async function updateKnowledgeBase( const [currentKb] = await tx .select({ workspaceId: knowledgeBase.workspaceId, userId: knowledgeBase.userId }) .from(knowledgeBase) - .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt))) + .where( + and( + eq(knowledgeBase.id, knowledgeBaseId), + isNull(knowledgeBase.deletedAt), + options?.assertedWorkspaceId + ? eq(knowledgeBase.workspaceId, options.assertedWorkspaceId) + : undefined + ) + ) .for('update') .limit(1) @@ -632,7 +786,15 @@ export async function updateKnowledgeBase( await tx .update(knowledgeBase) .set(updateData) - .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt))) + .where( + and( + eq(knowledgeBase.id, knowledgeBaseId), + isNull(knowledgeBase.deletedAt), + options?.assertedWorkspaceId + ? eq(knowledgeBase.workspaceId, options.assertedWorkspaceId) + : undefined + ) + ) // When a KB changes workspace, re-point the ownership bindings for its // stored files so file authorization (which resolves the owning workspace @@ -726,7 +888,15 @@ export async function updateKnowledgeBase( isNull(document.deletedAt) ) ) - .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt))) + .where( + and( + eq(knowledgeBase.id, knowledgeBaseId), + isNull(knowledgeBase.deletedAt), + options?.assertedWorkspaceId + ? eq(knowledgeBase.workspaceId, options.assertedWorkspaceId) + : undefined + ) + ) .groupBy(knowledgeBase.id) .limit(1) @@ -804,12 +974,26 @@ export async function getKnowledgeBaseById( export async function deleteKnowledgeBase( knowledgeBaseId: string, requestId: string, - options?: { archivedAt?: Date } + options?: { archivedAt?: Date; assertedWorkspaceId?: string } ): Promise { const now = options?.archivedAt ?? new Date() await db.transaction(async (tx) => { - await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`) + const [locked] = await tx + .select({ id: knowledgeBase.id, workspaceId: knowledgeBase.workspaceId }) + .from(knowledgeBase) + .where( + and( + eq(knowledgeBase.id, knowledgeBaseId), + isNull(knowledgeBase.deletedAt), + options?.assertedWorkspaceId + ? eq(knowledgeBase.workspaceId, options.assertedWorkspaceId) + : undefined + ) + ) + .limit(1) + .for('update') + if (!locked) throw new KnowledgeBaseNotFoundError(knowledgeBaseId) await tx .update(knowledgeBase) @@ -817,7 +1001,15 @@ export async function deleteKnowledgeBase( deletedAt: now, updatedAt: now, }) - .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt))) + .where( + and( + eq(knowledgeBase.id, knowledgeBaseId), + isNull(knowledgeBase.deletedAt), + options?.assertedWorkspaceId + ? eq(knowledgeBase.workspaceId, options.assertedWorkspaceId) + : undefined + ) + ) await tx .update(document) diff --git a/apps/sim/lib/uploads/upload-session/service.test.ts b/apps/sim/lib/uploads/upload-session/service.test.ts index ea5901633c2..020e3f2c6b1 100644 --- a/apps/sim/lib/uploads/upload-session/service.test.ts +++ b/apps/sim/lib/uploads/upload-session/service.test.ts @@ -1,6 +1,7 @@ /** * @vitest-environment node */ +import type { Principal } from '@sim/auth/principal' import { sha256Hex } from '@sim/security/hash' import { dbChainMockFns, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' import { eq, inArray, isNull } from 'drizzle-orm' @@ -60,8 +61,11 @@ import { assertUploadSessionAuthBinding, cleanupExpiredUploadSessions, completeUploadSession, + createUploadPartUrls, createUploadSession, + createUploadSessionAuthBinding, getOwnedUploadSession, + getPrincipalKnowledgeDocumentUploadSession, UPLOAD_SESSION_PART_SIZE, UPLOAD_SESSION_PUT_MAX_BYTES, type UploadSessionRecord, @@ -139,6 +143,103 @@ describe('upload sessions', () => { ).rejects.toMatchObject({ code: 'not_found' }) }) + it('binds new knowledge-document sessions to the exact creating credential', async () => { + const row = uploadRow({ + purpose: 'knowledge_document', + knowledgeBaseId: 'kb-1', + storageContext: 'knowledge-base', + finalKey: 'kb/guide.pdf', + fileName: 'guide.pdf', + contentType: 'application/pdf', + }) + dbChainMockFns.returning.mockResolvedValueOnce([row]) + + await createUploadSession({ + id: row.id, + workspaceId: WORKSPACE_ID, + knowledgeBaseId: 'kb-1', + userId: 'user-1', + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + purpose: 'knowledge_document', + fileName: 'guide.pdf', + contentType: 'application/pdf', + fileSize: 4, + localOrigin: 'http://localhost:3000', + }) + + expect(dbChainMockFns.values.mock.calls[0][0].metadata.authBinding).toEqual({ + version: 1, + workspaceId: WORKSPACE_ID, + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + }) + }) + + it('rejects a different API key on a bound knowledge-document control leg', async () => { + const row = uploadRow({ + purpose: 'knowledge_document', + knowledgeBaseId: 'kb-1', + storageContext: 'knowledge-base', + metadata: { + authBinding: { + version: 1, + workspaceId: WORKSPACE_ID, + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + }, + }, + }) + queueTableRows(schemaMock.uploadSession, [row]) + + await expect( + getPrincipalKnowledgeDocumentUploadSession({ + uploadId: row.id, + uploadToken: 'upload-secret', + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-2' }, + workspaceId: WORKSPACE_ID, + knowledgeBaseId: 'kb-1', + }) + ).rejects.toMatchObject({ code: 'not_found' }) + }) + + it.each([ + { + label: 'session', + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + mismatch: { kind: 'session', userId: 'user-1', sessionId: 'session-2' }, + }, + { + label: 'personal API key', + principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' }, + mismatch: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-2' }, + }, + { + label: 'workspace API key', + principal: { + kind: 'workspace_api_key', + workspaceId: WORKSPACE_ID, + keyId: 'workspace-key-1', + }, + mismatch: { + kind: 'workspace_api_key', + workspaceId: WORKSPACE_ID, + keyId: 'workspace-key-2', + }, + }, + ] satisfies Array<{ label: string; principal: Principal; mismatch: Principal }>)( + 'requires the exact bound $label credential for knowledge control', + ({ principal, mismatch }) => { + const session = sessionRecord({ + purpose: 'knowledge_document', + knowledgeBaseId: 'kb-1', + metadata: { authBinding: createUploadSessionAuthBinding(principal, WORKSPACE_ID) }, + }) + + expect(() => assertUploadSessionAuthBinding(session, principal)).not.toThrow() + expect(() => assertUploadSessionAuthBinding(session, mismatch)).toThrow( + 'Upload session not found' + ) + } + ) + it('preserves legacy unbound sessions under their prior ownership rules', () => { const legacy = sessionRecord({ metadata: {} }) @@ -180,8 +281,35 @@ describe('upload sessions', () => { ).toThrow('Upload session not found') }) + it('preserves the explicit missing-binding compatibility path for old knowledge sessions', () => { + const legacy = sessionRecord({ + purpose: 'knowledge_document', + knowledgeBaseId: 'kb-1', + metadata: {}, + }) + + expect(() => + assertUploadSessionAuthBinding(legacy, { + kind: 'personal_api_key', + userId: legacy.userId, + keyId: 'replacement-key', + }) + ).not.toThrow() + expect(() => + assertUploadSessionAuthBinding(legacy, { + kind: 'personal_api_key', + userId: 'different-user', + keyId: 'replacement-key', + }) + ).toThrow('Upload session not found') + }) + it('never treats a malformed credential binding as a legacy session', () => { - const malformed = sessionRecord({ metadata: { authBinding: { version: 1 } } }) + const malformed = sessionRecord({ + purpose: 'knowledge_document', + knowledgeBaseId: 'kb-1', + metadata: { authBinding: { version: 1 } }, + }) expect(() => assertUploadSessionAuthBinding(malformed, { @@ -216,6 +344,38 @@ describe('upload sessions', () => { expect(mockCreatePutTransfer).not.toHaveBeenCalled() }) + it('preserves multipart request bounds before provider signing', async () => { + const multipart = sessionRecord({ + method: 'multipart', + providerUploadId: 'provider-upload-1', + partSize: UPLOAD_SESSION_PART_SIZE, + partCount: 2, + fileSize: UPLOAD_SESSION_PART_SIZE + 1, + }) + + await expect( + createUploadPartUrls({ + session: multipart, + partNumbers: [1, 1], + localOrigin: 'http://localhost:3000', + }) + ).rejects.toMatchObject({ code: 'validation' }) + await expect( + createUploadPartUrls({ + session: multipart, + partNumbers: Array.from({ length: 101 }, (_, index) => index + 1), + localOrigin: 'http://localhost:3000', + }) + ).rejects.toMatchObject({ code: 'validation' }) + await expect( + createUploadPartUrls({ + session: multipart, + partNumbers: [3], + localOrigin: 'http://localhost:3000', + }) + ).rejects.toMatchObject({ code: 'validation' }) + }) + it('loads ownership from PostgreSQL and rejects a mismatched token', async () => { const token = 'upload-secret' const row = uploadRow({ tokenHash: sha256Hex(token) }) diff --git a/apps/sim/lib/uploads/upload-session/service.ts b/apps/sim/lib/uploads/upload-session/service.ts index fa6f44e20bf..b6e3c7303e9 100644 --- a/apps/sim/lib/uploads/upload-session/service.ts +++ b/apps/sim/lib/uploads/upload-session/service.ts @@ -92,7 +92,7 @@ export interface UploadSessionRecord { } /** - * The credential that was authorized to create a workspace-file upload. + * The credential that was authorized to create a protected upload session. * * This is deliberately kept in the existing JSON metadata column. It is * server-authored and immutable for the lifetime of the session; the upload @@ -130,13 +130,18 @@ interface CreateUploadSessionBaseParams { fileSize: number metadata?: Record localOrigin?: string - principal?: Principal } export type CreateUploadSessionParams = CreateUploadSessionBaseParams & ( - | { purpose: 'workspace_file' | 'table_import'; workspaceId: string } - | { purpose: 'knowledge_document'; workspaceId: string; knowledgeBaseId: string } + | { purpose: 'workspace_file'; workspaceId: string; principal: Principal } + | { purpose: 'table_import'; workspaceId: string } + | { + purpose: 'knowledge_document' + workspaceId: string + knowledgeBaseId: string + principal: Principal + } | { purpose: 'profile_picture'; workspaceId?: null } | { purpose: 'workspace_logo' | 'mothership_attachment'; workspaceId: string } | { @@ -157,10 +162,10 @@ export async function createUploadSession( const uploadToken = generateSecureToken(32) const workspaceId = params.purpose === 'profile_picture' ? null : params.workspaceId const metadata = { ...(params.metadata ?? {}) } - if (params.purpose === 'workspace_file') { - if (!workspaceId) throw new Error('Workspace-file upload is missing workspaceId') + if (params.purpose === 'workspace_file' || params.purpose === 'knowledge_document') { + if (!workspaceId) throw new Error(`${params.purpose} upload is missing workspaceId`) if (!params.principal) { - throw new Error('Workspace-file upload requires an authenticated principal') + throw new Error(`${params.purpose} upload requires an authenticated principal`) } metadata.authBinding = createUploadSessionAuthBinding(params.principal, workspaceId) } @@ -315,7 +320,7 @@ export async function getOwnedUploadSession(params: { if (params.executionId !== undefined && session.executionId !== params.executionId) { throw uploadNotFound() } - if (params.principal && session.purpose === 'workspace_file') { + if (params.principal && isPrincipalBoundUploadPurpose(session.purpose)) { assertUploadSessionAuthBinding(session, params.principal) } return session @@ -341,6 +346,24 @@ export async function getPrincipalUploadSession(params: { return session } +/** Loads a knowledge-document session using its immutable credential binding. */ +export async function getPrincipalKnowledgeDocumentUploadSession(params: { + uploadId: string + uploadToken: string + principal: Principal + workspaceId: string + knowledgeBaseId: string +}): Promise { + return getOwnedUploadSession({ + uploadId: params.uploadId, + uploadToken: params.uploadToken, + workspaceId: params.workspaceId, + purpose: 'knowledge_document', + knowledgeBaseId: params.knowledgeBaseId, + principal: params.principal, + }) +} + export function createUploadSessionAuthBinding( principal: Principal, workspaceId: string @@ -380,7 +403,7 @@ export function assertUploadSessionAuthBinding( session: UploadSessionRecord, principal: Principal ): void { - if (session.purpose !== 'workspace_file') return + if (!isPrincipalBoundUploadPurpose(session.purpose)) return const candidate = session.metadata.authBinding if (candidate === undefined) { assertLegacyUploadSessionOwner(session, principal) @@ -408,10 +431,10 @@ export function assertUploadSessionAuthBinding( /** * Preserves control access for the bounded set of sessions created before - * immutable credential bindings shipped. New workspace-file sessions always - * persist `authBinding`, and malformed bindings never enter this compatibility - * path. The upload token and current workspace authorization are still checked - * by the calling control-plane use case. + * immutable credential bindings shipped. New protected sessions always persist + * `authBinding`, and malformed bindings never enter this compatibility path. + * The upload token and current workspace authorization are still checked by the + * calling control-plane use case. */ function assertLegacyUploadSessionOwner(session: UploadSessionRecord, principal: Principal): void { const matches = @@ -1083,6 +1106,10 @@ function requiresStorageQuota(purpose: UploadSessionPurpose): boolean { return purpose === 'workspace_file' || purpose === 'knowledge_document' } +function isPrincipalBoundUploadPurpose(purpose: UploadSessionPurpose): boolean { + return purpose === 'workspace_file' || purpose === 'knowledge_document' +} + function resolveUploadStorage( params: CreateUploadSessionParams, id: string