guardrail attribution

Author: icecrasher321

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

@@ -275,6 +275,24 @@ describe('POST /api/files/multipart action=initiate quota enforcement', () => {
     expect(mockInitiateS3MultipartUpload).toHaveBeenCalled()
   })
 
+  it('keeps mothership chat uploads outside workspace storage quotas', async () => {
+    mockCheckStorageQuota.mockResolvedValue({ allowed: false, error: 'Storage limit exceeded' })
+
+    const res = await makeInitiateRequest({
+      fileName: 'conversation.bin',
+      contentType: 'application/octet-stream',
+      fileSize: 99999,
+      workspaceId: 'ws-1',
+      context: 'mothership',
+    })
+
+    const response = await POST(res)
+    expect(response.status).toBe(200)
+    expect(mockResolveStorageBillingContext).not.toHaveBeenCalled()
+    expect(mockCheckStorageQuota).not.toHaveBeenCalled()
+    expect(mockInitiateS3MultipartUpload).toHaveBeenCalled()
+  })
+
   it.each(['og-images', 'profile-pictures', 'workspace-logos', 'logs'])(
     'rejects quota-exempt context %s — not allowed via the multipart endpoint',
     async (context) => {

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

@@ -31,13 +31,11 @@ import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
 const logger = createLogger('MultipartUploadAPI')
 
 /**
- * Contexts the multipart endpoint accepts. The quota-exempt public-asset
- * contexts (`profile-pictures`, `workspace-logos`, `og-images`) and the
- * system-internal `logs` context are deliberately excluded: their uploads are
- * small images capped far below the multipart threshold and routed through the
- * presigned endpoint, so they have no large-file flow here. Accepting them would
- * only expose a path that bypasses the per-user storage quota, since every
- * context in this set is quota-enforced below.
+ * Contexts the multipart endpoint accepts. Small public assets and internal logs
+ * are excluded because they have no large-file flow. Mothership remains
+ * available for large chat attachments but is quota-exempt because chat uploads
+ * do not count as durable workspace-file storage. Every other accepted context
+ * is quota-enforced below.
  */
 const ALLOWED_UPLOAD_CONTEXTS = new Set<StorageContext>([
   'knowledge-base',

apps/sim/app/api/guardrails/validate/route.test.ts

@@ -266,4 +266,47 @@ describe('POST /api/guardrails/validate', () => {
       })
     )
   })
+
+  it('rejects invalid internal billing attribution as a protocol error', async () => {
+    hybridAuthMockFns.mockCheckSessionOrInternalAuth.mockResolvedValue({
+      success: true,
+      userId: 'user-1',
+      authType: 'internal_jwt',
+    })
+    mockRequireBillingAttributionHeader.mockImplementationOnce(() => {
+      throw new Error('Billing attribution header is required')
+    })
+
+    const res = await POST(
+      createMockRequest('POST', {
+        validationType: 'hallucination',
+        input: 'test input',
+        knowledgeBaseId: 'kb-1',
+        model: 'gpt-4o',
+        workflowId: 'wf-1',
+      })
+    )
+
+    expect(res.status).toBe(400)
+    await expect(res.json()).resolves.toEqual({ error: 'Invalid billing attribution' })
+    expect(mockValidateHallucination).not.toHaveBeenCalled()
+  })
+
+  it('surfaces workspace billing resolution failures as infrastructure errors', async () => {
+    mockResolveBillingAttribution.mockRejectedValueOnce(new Error('Database unavailable'))
+
+    const res = await POST(
+      createMockRequest('POST', {
+        validationType: 'hallucination',
+        input: 'test input',
+        knowledgeBaseId: 'kb-1',
+        model: 'gpt-4o',
+        workflowId: 'wf-1',
+      })
+    )
+
+    expect(res.status).toBe(500)
+    await expect(res.json()).resolves.toEqual({ error: 'Failed to resolve billing attribution' })
+    expect(mockValidateHallucination).not.toHaveBeenCalled()
+  })
 })

apps/sim/app/api/guardrails/validate/route.ts

@@ -166,16 +166,29 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
       }
 
       resolvedWorkspaceId = authorization.workflow.workspaceId
-      billingAttribution =
-        auth.authType === AuthType.INTERNAL_JWT
-          ? requireBillingAttributionHeader(request.headers, {
-              actorUserId: auth.userId,
-              workspaceId: resolvedWorkspaceId,
-            })
-          : await resolveBillingAttribution({
-              actorUserId: auth.userId,
-              workspaceId: resolvedWorkspaceId,
-            })
+      try {
+        billingAttribution =
+          auth.authType === AuthType.INTERNAL_JWT
+            ? requireBillingAttributionHeader(request.headers, {
+                actorUserId: auth.userId,
+                workspaceId: resolvedWorkspaceId,
+              })
+            : await resolveBillingAttribution({
+                actorUserId: auth.userId,
+                workspaceId: resolvedWorkspaceId,
+              })
+      } catch (error) {
+        const isInternalRequest = auth.authType === AuthType.INTERNAL_JWT
+        logger.error(`[${requestId}] Failed to establish billing attribution`, { error })
+        return NextResponse.json(
+          {
+            error: isInternalRequest
+              ? 'Invalid billing attribution'
+              : 'Failed to resolve billing attribution',
+          },
+          { status: isInternalRequest ? 400 : 500 }
+        )
+      }
 
       try {
         await assertPermissionsAllowed({

apps/sim/lib/uploads/shared/types.ts

@@ -26,16 +26,16 @@ export type StorageContext =
 /**
  * Contexts exempt from storage quota checks. Includes system-internal contexts
  * (`logs` — written by the execution pipeline, not user-initiated) and small
- * metadata assets (`profile-pictures`, `workspace-logos`, `og-images`). All
- * other contexts are user-driven uploads and must pass quota validation.
+ * metadata assets (`profile-pictures`, `workspace-logos`, `og-images`).
+ * Mothership chat attachments are also exempt because they are not counted as
+ * durable workspace-file storage.
  *
- * Note: every quota-exempt context is excluded from `ALLOWED_UPLOAD_CONTEXTS`
- * in the multipart endpoint, so none are reachable there — the exemption applies
- * only to the single-part upload paths (presigned/FormData) those small assets
- * actually use. The multipart endpoint therefore only ever serves
- * quota-enforced contexts.
+ * The small-asset and system contexts are excluded from the multipart endpoint.
+ * Mothership remains available there for large chat attachments while retaining
+ * the same quota exemption as its single-part upload path.
  */
 export const QUOTA_EXEMPT_STORAGE_CONTEXTS = new Set<StorageContext>([
+  'mothership',
   'profile-pictures',
   'workspace-logos',
   'og-images',