Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions apps/api/src/app/s3.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import {
GetObjectCommand,
HeadObjectCommand,
PutObjectCommand,
S3Client,
type GetObjectCommandOutput,
Expand Down Expand Up @@ -195,3 +196,27 @@ export async function getObjectAsBuffer(
const bytes = await response.Body.transformToByteArray();
return Buffer.from(bytes);
}

/**
* Fetch an S3 object's size (in bytes) via a HEAD request, WITHOUT downloading
* the body. Used to reject oversized uploads before loading them into memory —
* `getObjectAsBuffer` would otherwise buffer the entire object (and base64
* callers expand it ~1.33x on top), so a single huge file could OOM the API.
*
* Returns `undefined` if S3 doesn't report a ContentLength (callers should treat
* that as "size unknown" rather than "zero").
*/
export async function getObjectContentLength(
bucket: string,
key: string,
): Promise<number | undefined> {
if (!s3Client) {
throw new Error('S3 client not configured');
}

const response = await s3Client.send(
new HeadObjectCommand({ Bucket: bucket, Key: key }),
);

return response.ContentLength;
}
4 changes: 3 additions & 1 deletion apps/api/src/attachments/attachments.module.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import { Module } from '@nestjs/common';
import { AuthModule } from '../auth/auth.module';
import { UploadsModule } from '../uploads/uploads.module';
import { AttachmentsController } from './attachments.controller';
import { AttachmentsService } from './attachments.service';

@Module({
imports: [AuthModule], // Import AuthModule for HybridAuthGuard dependencies
// AuthModule: HybridAuthGuard deps. UploadsModule: presigned-upload s3Key reads.
imports: [AuthModule, UploadsModule],
controllers: [AttachmentsController],
providers: [AttachmentsService],
exports: [AttachmentsService],
Expand Down
86 changes: 86 additions & 0 deletions apps/api/src/attachments/attachments.service.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
import { BadRequestException } from '@nestjs/common';

// Mocks must be declared before importing the service under test.
jest.mock('@/app/s3', () => ({
s3Client: { send: jest.fn().mockResolvedValue({}) },
getSignedUrl: jest.fn().mockResolvedValue('https://signed.example/file'),
}));

jest.mock('@db', () => ({
db: { attachment: { create: jest.fn() } },
AttachmentType: {
image: 'image',
video: 'video',
audio: 'audio',
document: 'document',
other: 'other',
},
AttachmentEntityType: { task: 'task', offboarding_checklist: 'offboarding_checklist' },
}));

jest.mock('../utils/file-type-validation', () => ({
validateFileContent: jest.fn(),
}));

import { db } from '@db';
import { AttachmentsService } from './attachments.service';

const mockUploadsService = { readUploadAsBase64: jest.fn() };

describe('AttachmentsService — presigned s3Key uploads', () => {
let service: AttachmentsService;

beforeEach(() => {
jest.clearAllMocks();
process.env.APP_AWS_BUCKET_NAME = 'test-bucket';
service = new AttachmentsService(mockUploadsService as never);
});

it('resolves the file from s3Key (presigned) — no base64 through the LLM — and uploads it', async () => {
mockUploadsService.readUploadAsBase64.mockResolvedValue(
Buffer.from('hello world').toString('base64'),
);
(db.attachment.create as jest.Mock).mockResolvedValue({
id: 'att_1',
name: 'rbac.pdf',
type: 'document',
url: 'org_1/attachments/task/tsk_1/key',
createdAt: new Date(),
});

const result = await service.uploadAttachment(
'org_1',
'tsk_1',
'task' as never,
{
fileName: 'rbac.pdf',
fileType: 'application/pdf',
s3Key: 'org_1/uploads/attachment/123-rbac.pdf',
} as never,
'usr_1',
);

// Fetched the bytes from the org-scoped presigned key instead of base64.
expect(mockUploadsService.readUploadAsBase64).toHaveBeenCalledWith(
'org_1',
'org_1/uploads/attachment/123-rbac.pdf',
);
expect(db.attachment.create).toHaveBeenCalled();
expect(result.id).toBe('att_1');
});

it('throws when neither fileData nor s3Key is provided', async () => {
await expect(
service.uploadAttachment(
'org_1',
'tsk_1',
'task' as never,
{ fileName: 'rbac.pdf', fileType: 'application/pdf' } as never,
'usr_1',
),
).rejects.toBeInstanceOf(BadRequestException);

expect(mockUploadsService.readUploadAsBase64).not.toHaveBeenCalled();
expect(db.attachment.create).not.toHaveBeenCalled();
});
});
23 changes: 21 additions & 2 deletions apps/api/src/attachments/attachments.service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import {
import { randomBytes } from 'crypto';
import { AttachmentResponseDto } from '../tasks/dto/task-responses.dto';
import { UploadAttachmentDto } from './upload-attachment.dto';
import { UploadsService } from '../uploads/uploads.service';
import { validateFileContent } from '../utils/file-type-validation';

@Injectable()
Expand All @@ -24,7 +25,7 @@ export class AttachmentsService {
private readonly MAX_FILE_SIZE_BYTES = 100 * 1024 * 1024; // 100MB
private readonly SIGNED_URL_EXPIRY = 900; // 15 minutes

constructor() {
constructor(private readonly uploadsService: UploadsService) {
// AWS configuration is validated at startup via ConfigModule
// Safe to access environment variables directly since they're validated
this.bucketName = process.env.APP_AWS_BUCKET_NAME!;
Expand Down Expand Up @@ -115,8 +116,26 @@ export class AttachmentsService {
);
}

// Resolve the file content from either inline base64 (UI/direct callers)
// or a presigned-upload s3Key (AI/MCP clients — avoids slow base64 through
// an LLM). readUploadAsBase64 enforces that the key belongs to this org.
const fileData =
uploadDto.fileData ??
(uploadDto.s3Key
? await this.uploadsService.readUploadAsBase64(
organizationId,
uploadDto.s3Key,
)
: undefined);

if (!fileData) {
throw new BadRequestException(
'Provide either fileData (base64) or s3Key from /v1/uploads/presign.',
);
}

// Validate file size
const fileBuffer = Buffer.from(uploadDto.fileData, 'base64');
const fileBuffer = Buffer.from(fileData, 'base64');
if (fileBuffer.length > this.MAX_FILE_SIZE_BYTES) {
throw new BadRequestException(
`File size exceeds maximum allowed size of ${this.MAX_FILE_SIZE_BYTES / (1024 * 1024)}MB`,
Expand Down
21 changes: 18 additions & 3 deletions apps/api/src/attachments/upload-attachment.dto.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import {
MaxLength,
} from 'class-validator';
import { IsMimeTypeField } from '../utils/mime-type.validator';
import { MAX_UPLOAD_BASE64_LENGTH } from '../uploads/upload-limits';

export class UploadAttachmentDto {
@ApiProperty({
Expand All @@ -29,15 +30,29 @@ export class UploadAttachmentDto {
fileType: string;

@ApiProperty({
description: 'Base64 encoded file data',
description:
'Base64-encoded file contents. For the web UI / direct callers. AI/MCP clients should instead upload via /v1/uploads/presign (purpose=attachment) and pass `s3Key` — base64 through an LLM is impractically slow and times out. Provide exactly one of fileData or s3Key.',
required: false,
example:
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==',
})
@IsOptional()
@IsString()
@IsNotEmpty()
@MaxLength(134_217_728)
@MaxLength(MAX_UPLOAD_BASE64_LENGTH)
@IsBase64()
fileData: string;
fileData?: string;

@ApiProperty({
description:
'Key of a file already uploaded via /v1/uploads/presign (purpose=attachment). The server fetches the bytes from storage — no base64 needed. Provide exactly one of fileData or s3Key.',
required: false,
example: 'org_abc123/uploads/attachment/1700000000000-rbac-matrix.xlsx',
})
@IsOptional()
@IsString()
@IsNotEmpty()
s3Key?: string;

@ApiProperty({
description: 'Description of the attachment',
Expand Down
34 changes: 32 additions & 2 deletions apps/api/src/knowledge-base/dto/upload-document.dto.ts
Original file line number Diff line number Diff line change
@@ -1,18 +1,48 @@
import { IsOptional, IsString } from 'class-validator';
import { ApiProperty } from '@nestjs/swagger';
import { IsBase64, IsOptional, IsString, MaxLength } from 'class-validator';
import { MAX_UPLOAD_BASE64_LENGTH } from '../../uploads/upload-limits';

export class UploadDocumentDto {
@ApiProperty({ description: 'Organization ID that owns the document' })
@IsString()
organizationId!: string;

@ApiProperty({ description: 'File name', example: 'rbac-matrix.xlsx' })
@IsString()
fileName!: string;

@ApiProperty({
description: 'MIME type of the file',
example:
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
})
@IsString()
fileType!: string;

@ApiProperty({
description:
'Base64-encoded file contents. For the web UI / direct callers. AI/MCP clients should instead upload via /v1/uploads/presign (purpose=document) and pass `s3Key` — base64 through an LLM is impractically slow and times out. Provide exactly one of fileData or s3Key.',
required: false,
})
@IsOptional()
@IsString()
// Cap the inline payload at the validation layer (before it is decoded),
// matching the other migrated upload DTOs. The limit is the base64 length of
// the 100 MiB file ceiling — see upload-limits.ts.
@MaxLength(MAX_UPLOAD_BASE64_LENGTH)
@IsBase64()
fileData?: string; // base64 encoded

@ApiProperty({
description:
'Key of a file already uploaded via /v1/uploads/presign (purpose=document). The server fetches the bytes from storage — no base64 needed. Provide exactly one of fileData or s3Key.',
required: false,
})
@IsOptional()
@IsString()
fileData!: string; // base64 encoded
s3Key?: string;

@ApiProperty({ description: 'Optional description', required: false })
@IsOptional()
@IsString()
description?: string;
Expand Down
3 changes: 2 additions & 1 deletion apps/api/src/knowledge-base/knowledge-base.module.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import { Module } from '@nestjs/common';
import { AuthModule } from '../auth/auth.module';
import { UploadsModule } from '../uploads/uploads.module';
import { KnowledgeBaseController } from './knowledge-base.controller';
import { KnowledgeBaseService } from './knowledge-base.service';

@Module({
imports: [AuthModule],
imports: [AuthModule, UploadsModule],
controllers: [KnowledgeBaseController],
providers: [KnowledgeBaseService],
})
Expand Down
54 changes: 53 additions & 1 deletion apps/api/src/knowledge-base/knowledge-base.service.spec.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,11 @@
import { BadRequestException } from '@nestjs/common';
import { Test, TestingModule } from '@nestjs/testing';
import { KnowledgeBaseService } from './knowledge-base.service';
import { UploadsService } from '../uploads/uploads.service';

const mockUploadsService = {
readUploadAsBase64: jest.fn(),
};

jest.mock('@db', () => ({
db: {
Expand Down Expand Up @@ -68,7 +74,10 @@ describe('KnowledgeBaseService', () => {

beforeEach(async () => {
const module: TestingModule = await Test.createTestingModule({
providers: [KnowledgeBaseService],
providers: [
KnowledgeBaseService,
{ provide: UploadsService, useValue: mockUploadsService },
],
}).compile();

service = module.get<KnowledgeBaseService>(KnowledgeBaseService);
Expand Down Expand Up @@ -215,6 +224,49 @@ describe('KnowledgeBaseService', () => {
'base64data',
);
});

it('resolves content from s3Key (presigned upload) when no fileData', async () => {
mockUploadsService.readUploadAsBase64.mockResolvedValue('fromS3base64');
(uploadToS3 as jest.Mock).mockResolvedValue({
s3Key: 'org_1/doc.pdf',
fileSize: 2048,
});
(mockDb.knowledgeBaseDocument.create as jest.Mock).mockResolvedValue({
id: 'd2',
name: 'doc.pdf',
s3Key: 'org_1/doc.pdf',
});

await service.uploadDocument({
organizationId: 'org_1',
fileName: 'doc.pdf',
fileType: 'application/pdf',
s3Key: 'org_1/uploads/document/123-doc.pdf',
} as any);

// Fetched the bytes from the presigned key (org-scoped, no base64 via LLM)
expect(mockUploadsService.readUploadAsBase64).toHaveBeenCalledWith(
'org_1',
'org_1/uploads/document/123-doc.pdf',
);
expect(uploadToS3).toHaveBeenCalledWith(
'org_1',
'doc.pdf',
'application/pdf',
'fromS3base64',
);
});

it('throws when neither fileData nor s3Key is provided', async () => {
await expect(
service.uploadDocument({
organizationId: 'org_1',
fileName: 'doc.pdf',
fileType: 'application/pdf',
} as any),
).rejects.toBeInstanceOf(BadRequestException);
expect(uploadToS3).not.toHaveBeenCalled();
});
});

describe('getDownloadUrl', () => {
Expand Down
Loading
Loading