mirror of
https://github.com/lobehub/lobehub
synced 2026-04-21 17:47:27 +00:00
fix(image): not show reference image (#11338)
This commit is contained in:
YuTengjing
GitHub
parent
815596de02
commit
de1504cf7b
18 changed files with 704 additions and 355 deletions
|
|
@ -52,3 +52,4 @@ alwaysApply: false
|
|||
|
||||
- Never log user private information like api key, etc
|
||||
- Don't use `import { log } from 'debug'` to log messages, because it will directly log the message to the console.
|
||||
- Use console.error instead of debug package to log error message in catch block.
|
||||
|
|
|
|||
|
|
@ -77,8 +77,8 @@ const GenerationFeed = memo(() => {
|
|||
}
|
||||
|
||||
return (
|
||||
<Flexbox flex={1}>
|
||||
<Flexbox gap={16} ref={parent} width="100%">
|
||||
<Flexbox flex={1} style={{ overflowY: 'auto' }}>
|
||||
<Flexbox gap={16} ref={parent} style={{ paddingBottom: 48 }} width="100%">
|
||||
{currentGenerationBatches.map((batch, index) => (
|
||||
<Fragment key={batch.id}>
|
||||
{Boolean(index !== 0) && <Divider dashed style={{ margin: 0 }} />}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,5 @@
|
|||
'use client';
|
||||
|
||||
import { Center } from '@lobehub/ui';
|
||||
|
||||
import { useImageStore } from '@/store/image';
|
||||
import { generationBatchSelectors, generationTopicSelectors } from '@/store/image/selectors';
|
||||
|
||||
|
|
@ -32,15 +30,7 @@ const ImageWorkspaceContent = () => {
|
|||
<GenerationFeed key={activeTopicId} />
|
||||
|
||||
{/* 底部输入框 */}
|
||||
<Center
|
||||
style={{
|
||||
bottom: 24,
|
||||
position: 'sticky',
|
||||
width: '100%',
|
||||
}}
|
||||
>
|
||||
<PromptInput disableAnimation={true} showTitle={false} />
|
||||
</Center>
|
||||
<PromptInput disableAnimation={true} showTitle={false} />
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -120,13 +120,7 @@ const PromptInput = ({ showTitle = false }: PromptInputProps) => {
|
|||
};
|
||||
|
||||
return (
|
||||
<Flexbox
|
||||
gap={32}
|
||||
style={{
|
||||
marginTop: 48,
|
||||
}}
|
||||
width={'100%'}
|
||||
>
|
||||
<Flexbox gap={32} width={'100%'}>
|
||||
{showTitle && <PromptTitle />}
|
||||
<ChatInput
|
||||
className={cx(styles.container, isDarkMode && styles.container_dark)}
|
||||
|
|
|
|||
|
|
@ -1,14 +1,13 @@
|
|||
'use client';
|
||||
|
||||
import { Flexbox } from '@lobehub/ui';
|
||||
import { Suspense, memo } from 'react';
|
||||
import { memo } from 'react';
|
||||
|
||||
import NavHeader from '@/features/NavHeader';
|
||||
import WideScreenContainer from '@/features/WideScreenContainer';
|
||||
import WideScreenButton from '@/features/WideScreenContainer/WideScreenButton';
|
||||
|
||||
import ImageWorkspace from './features/ImageWorkspace';
|
||||
import SkeletonList from './features/ImageWorkspace/SkeletonList';
|
||||
|
||||
const DesktopImagePage = memo(() => {
|
||||
return (
|
||||
|
|
@ -16,9 +15,7 @@ const DesktopImagePage = memo(() => {
|
|||
<NavHeader right={<WideScreenButton />} />
|
||||
<Flexbox height={'100%'} style={{ overflowY: 'auto', position: 'relative' }} width={'100%'}>
|
||||
<WideScreenContainer height={'100%'} wrapperStyle={{ height: '100%' }}>
|
||||
<Suspense fallback={<SkeletonList />}>
|
||||
<ImageWorkspace />
|
||||
</Suspense>
|
||||
<ImageWorkspace />
|
||||
</WideScreenContainer>
|
||||
</Flexbox>
|
||||
</>
|
||||
|
|
|
|||
|
|
@ -185,7 +185,7 @@ const categorizeError = (
|
|||
}
|
||||
|
||||
return {
|
||||
errorMessage: error.message || AsyncTaskErrorType.ServerError,
|
||||
errorMessage: error.message || error.error?.message || AsyncTaskErrorType.ServerError,
|
||||
errorType: AsyncTaskErrorType.ServerError,
|
||||
};
|
||||
};
|
||||
|
|
|
|||
491
src/server/routers/lambda/image/index.test.ts
Normal file
491
src/server/routers/lambda/image/index.test.ts
Normal file
|
|
@ -0,0 +1,491 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { AsyncTaskStatus, AsyncTaskType } from '@/types/asyncTask';
|
||||
|
||||
// Use vi.hoisted for variables used in vi.mock factory
|
||||
const {
|
||||
mockServerDB,
|
||||
mockGetKeyFromFullUrl,
|
||||
mockGetFullFileUrl,
|
||||
mockAsyncTaskModelUpdate,
|
||||
mockChargeBeforeGenerate,
|
||||
mockCreateAsyncCaller,
|
||||
} = vi.hoisted(() => ({
|
||||
mockServerDB: {
|
||||
transaction: vi.fn(),
|
||||
},
|
||||
mockGetKeyFromFullUrl: vi.fn(),
|
||||
mockGetFullFileUrl: vi.fn(),
|
||||
mockAsyncTaskModelUpdate: vi.fn(),
|
||||
mockChargeBeforeGenerate: vi.fn(),
|
||||
mockCreateAsyncCaller: vi.fn(),
|
||||
}));
|
||||
|
||||
// Mock debug
|
||||
vi.mock('debug', () => ({
|
||||
default: () => () => {},
|
||||
}));
|
||||
|
||||
// Mock auth related
|
||||
vi.mock('@lobechat/utils/server', () => ({
|
||||
getXorPayload: vi.fn(() => ({})),
|
||||
}));
|
||||
|
||||
// Mock database adaptor
|
||||
vi.mock('@/database/core/db-adaptor', () => ({
|
||||
getServerDB: vi.fn(async () => mockServerDB),
|
||||
}));
|
||||
|
||||
// Mock FileService
|
||||
vi.mock('@/server/services/file', () => ({
|
||||
FileService: vi.fn(() => ({
|
||||
getKeyFromFullUrl: mockGetKeyFromFullUrl,
|
||||
getFullFileUrl: mockGetFullFileUrl,
|
||||
})),
|
||||
}));
|
||||
|
||||
// Mock AsyncTaskModel
|
||||
vi.mock('@/database/models/asyncTask', () => ({
|
||||
AsyncTaskModel: vi.fn(() => ({
|
||||
update: mockAsyncTaskModelUpdate,
|
||||
})),
|
||||
}));
|
||||
|
||||
// Mock chargeBeforeGenerate
|
||||
vi.mock('@/business/server/image-generation/chargeBeforeGenerate', () => ({
|
||||
chargeBeforeGenerate: (params: any) => mockChargeBeforeGenerate(params),
|
||||
}));
|
||||
|
||||
// Mock async caller
|
||||
vi.mock('@/server/routers/async/caller', () => ({
|
||||
createAsyncCaller: mockCreateAsyncCaller,
|
||||
}));
|
||||
|
||||
// Mock drizzle-orm
|
||||
vi.mock('drizzle-orm', () => ({
|
||||
and: vi.fn((...args) => args),
|
||||
eq: vi.fn((a, b) => ({ a, b })),
|
||||
}));
|
||||
|
||||
// Mock database schemas
|
||||
vi.mock('@/database/schemas', () => ({
|
||||
asyncTasks: { id: 'asyncTasks.id', userId: 'asyncTasks.userId' },
|
||||
generationBatches: { id: 'generationBatches.id' },
|
||||
generations: { id: 'generations.id', userId: 'generations.userId' },
|
||||
}));
|
||||
|
||||
// Mock seed generator
|
||||
vi.mock('@/utils/number', () => ({
|
||||
generateUniqueSeeds: vi.fn((count: number) => Array.from({ length: count }, (_, i) => 1000 + i)),
|
||||
}));
|
||||
|
||||
import { imageRouter } from './index';
|
||||
|
||||
describe('imageRouter', () => {
|
||||
const mockUserId = 'test-user-id';
|
||||
const mockAsyncCallerCreateImage = vi.fn();
|
||||
|
||||
const createMockCtx = (overrides = {}) => ({
|
||||
userId: mockUserId,
|
||||
authorizationHeader: 'mock-auth-header',
|
||||
...overrides,
|
||||
});
|
||||
|
||||
const createDefaultInput = (overrides = {}) => ({
|
||||
generationTopicId: 'topic-1',
|
||||
imageNum: 2,
|
||||
model: 'stable-diffusion',
|
||||
params: {
|
||||
prompt: 'a beautiful sunset',
|
||||
width: 512,
|
||||
height: 512,
|
||||
},
|
||||
provider: 'test-provider',
|
||||
...overrides,
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
// Default mock implementations
|
||||
mockChargeBeforeGenerate.mockResolvedValue(undefined);
|
||||
mockGetKeyFromFullUrl.mockResolvedValue(null);
|
||||
mockGetFullFileUrl.mockResolvedValue(null);
|
||||
|
||||
// Setup default transaction mock
|
||||
const mockBatch = {
|
||||
id: 'batch-1',
|
||||
generationTopicId: 'topic-1',
|
||||
model: 'stable-diffusion',
|
||||
provider: 'test-provider',
|
||||
config: {},
|
||||
userId: mockUserId,
|
||||
};
|
||||
|
||||
const mockGenerations = [
|
||||
{ id: 'gen-1', generationBatchId: 'batch-1', seed: 1000, userId: mockUserId },
|
||||
{ id: 'gen-2', generationBatchId: 'batch-1', seed: 1001, userId: mockUserId },
|
||||
];
|
||||
|
||||
const mockAsyncTasks = [
|
||||
{ id: 'task-1', status: AsyncTaskStatus.Pending, type: AsyncTaskType.ImageGeneration },
|
||||
{ id: 'task-2', status: AsyncTaskStatus.Pending, type: AsyncTaskType.ImageGeneration },
|
||||
];
|
||||
|
||||
let insertCallCount = 0;
|
||||
mockServerDB.transaction.mockImplementation(async (callback) => {
|
||||
insertCallCount = 0;
|
||||
const tx = {
|
||||
insert: vi.fn().mockReturnValue({
|
||||
values: vi.fn().mockReturnValue({
|
||||
returning: vi.fn().mockImplementation(() => {
|
||||
insertCallCount++;
|
||||
if (insertCallCount === 1) return [mockBatch];
|
||||
if (insertCallCount === 2) return mockGenerations;
|
||||
// For async tasks, return one at a time
|
||||
const taskIndex = insertCallCount - 3;
|
||||
return [mockAsyncTasks[taskIndex] || mockAsyncTasks[0]];
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
update: vi.fn().mockReturnValue({
|
||||
set: vi.fn().mockReturnValue({
|
||||
where: vi.fn().mockResolvedValue(undefined),
|
||||
}),
|
||||
}),
|
||||
};
|
||||
|
||||
return callback(tx);
|
||||
});
|
||||
|
||||
mockCreateAsyncCaller.mockResolvedValue({
|
||||
image: {
|
||||
createImage: mockAsyncCallerCreateImage,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
describe('createImage', () => {
|
||||
it('should create image generation batch and generations successfully', async () => {
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput();
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.data.batch).toBeDefined();
|
||||
expect(result.data.batch.id).toBe('batch-1');
|
||||
expect(result.data.generations).toHaveLength(2);
|
||||
expect(mockServerDB.transaction).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should convert imageUrls to S3 keys for database storage', async () => {
|
||||
mockGetKeyFromFullUrl
|
||||
.mockResolvedValueOnce('files/image1.jpg')
|
||||
.mockResolvedValueOnce('files/image2.jpg');
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrls: [
|
||||
'https://s3.amazonaws.com/bucket/files/image1.jpg',
|
||||
'https://s3.amazonaws.com/bucket/files/image2.jpg',
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
await caller.createImage(input);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledTimes(2);
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(
|
||||
'https://s3.amazonaws.com/bucket/files/image1.jpg',
|
||||
);
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(
|
||||
'https://s3.amazonaws.com/bucket/files/image2.jpg',
|
||||
);
|
||||
});
|
||||
|
||||
it('should convert single imageUrl to S3 key for database storage', async () => {
|
||||
mockGetKeyFromFullUrl.mockResolvedValue('files/single-image.jpg');
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrl: 'https://s3.amazonaws.com/bucket/files/single-image.jpg',
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
await caller.createImage(input);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(
|
||||
'https://s3.amazonaws.com/bucket/files/single-image.jpg',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle failed URL to key conversion gracefully for imageUrls', async () => {
|
||||
mockGetKeyFromFullUrl.mockResolvedValue(null);
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrls: ['https://example.com/image.jpg'],
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
// Should still succeed, just with empty imageUrls in config
|
||||
expect(result.success).toBe(true);
|
||||
});
|
||||
|
||||
it('should throw error when imageUrls conversion fails and URLs remain', async () => {
|
||||
mockGetKeyFromFullUrl.mockRejectedValue(new Error('Conversion failed'));
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrls: ['https://example.com/image.jpg'],
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
|
||||
// When conversion fails, the original URL is kept but validateNoUrlsInConfig
|
||||
// will detect it and throw an error to prevent storing URLs in database
|
||||
await expect(caller.createImage(input)).rejects.toThrow(
|
||||
'Invalid configuration: Found full URL instead of key',
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw error when single imageUrl conversion fails and URL remains', async () => {
|
||||
mockGetKeyFromFullUrl.mockRejectedValue(new Error('Conversion failed'));
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrl: 'https://example.com/image.jpg',
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
|
||||
// When conversion fails, the original URL is kept but validateNoUrlsInConfig
|
||||
// will detect it and throw an error to prevent storing URLs in database
|
||||
await expect(caller.createImage(input)).rejects.toThrow(
|
||||
'Invalid configuration: Found full URL instead of key',
|
||||
);
|
||||
});
|
||||
|
||||
it('should return charge result when chargeBeforeGenerate returns a value', async () => {
|
||||
const chargeResult = {
|
||||
success: true as const,
|
||||
data: {
|
||||
batch: { id: 'charged-batch' },
|
||||
generations: [{ id: 'charged-gen' }],
|
||||
},
|
||||
};
|
||||
mockChargeBeforeGenerate.mockResolvedValue(chargeResult);
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput();
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
expect(result).toEqual(chargeResult);
|
||||
// Should not proceed with database transaction
|
||||
expect(mockServerDB.transaction).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should call chargeBeforeGenerate with correct parameters', async () => {
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput();
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
await caller.createImage(input);
|
||||
|
||||
expect(mockChargeBeforeGenerate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
generationTopicId: 'topic-1',
|
||||
imageNum: 2,
|
||||
model: 'stable-diffusion',
|
||||
provider: 'test-provider',
|
||||
userId: mockUserId,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should trigger async image generation tasks', async () => {
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput();
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
await caller.createImage(input);
|
||||
|
||||
expect(mockCreateAsyncCaller).toHaveBeenCalledWith({ userId: mockUserId });
|
||||
});
|
||||
|
||||
it('should handle async caller creation failure', async () => {
|
||||
mockCreateAsyncCaller.mockRejectedValue(new Error('Caller creation failed'));
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput();
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
// Should still return success as the database records were created
|
||||
expect(result.success).toBe(true);
|
||||
// Should update async task status to error
|
||||
expect(mockAsyncTaskModelUpdate).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should update all task statuses to error when async processing fails', async () => {
|
||||
mockCreateAsyncCaller.mockRejectedValue(new Error('Processing failed'));
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput();
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
await caller.createImage(input);
|
||||
|
||||
// Should update both tasks to error status
|
||||
expect(mockAsyncTaskModelUpdate).toHaveBeenCalledTimes(2);
|
||||
expect(mockAsyncTaskModelUpdate).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
expect.objectContaining({
|
||||
status: AsyncTaskStatus.Error,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should generate unique seeds when seed param is provided', async () => {
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
seed: 42,
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
await caller.createImage(input);
|
||||
|
||||
expect(mockServerDB.transaction).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should use null seeds when seed param is not provided', async () => {
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
// No seed param
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
await caller.createImage(input);
|
||||
|
||||
expect(mockServerDB.transaction).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should pass with valid key-based imageUrls', async () => {
|
||||
mockGetKeyFromFullUrl.mockResolvedValue('files/valid-key.jpg');
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrls: ['files/valid-key.jpg'],
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
});
|
||||
|
||||
describe('development environment URL conversion', () => {
|
||||
beforeEach(() => {
|
||||
vi.stubEnv('NODE_ENV', 'development');
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.unstubAllEnvs();
|
||||
});
|
||||
|
||||
it('should convert single imageUrl to S3 URL in development mode', async () => {
|
||||
mockGetKeyFromFullUrl.mockResolvedValue('files/image-key.jpg');
|
||||
mockGetFullFileUrl.mockResolvedValue('https://s3.amazonaws.com/bucket/files/image-key.jpg');
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrl: 'http://localhost:3000/f/file-id',
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(mockGetFullFileUrl).toHaveBeenCalledWith('files/image-key.jpg');
|
||||
});
|
||||
|
||||
it('should convert multiple imageUrls to S3 URLs in development mode', async () => {
|
||||
mockGetKeyFromFullUrl
|
||||
.mockResolvedValueOnce('files/image1.jpg')
|
||||
.mockResolvedValueOnce('files/image2.jpg');
|
||||
mockGetFullFileUrl
|
||||
.mockResolvedValueOnce('https://s3.amazonaws.com/bucket/files/image1.jpg')
|
||||
.mockResolvedValueOnce('https://s3.amazonaws.com/bucket/files/image2.jpg');
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrls: ['http://localhost:3000/f/id1', 'http://localhost:3000/f/id2'],
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(mockGetFullFileUrl).toHaveBeenCalledTimes(2);
|
||||
expect(mockGetFullFileUrl).toHaveBeenCalledWith('files/image1.jpg');
|
||||
expect(mockGetFullFileUrl).toHaveBeenCalledWith('files/image2.jpg');
|
||||
});
|
||||
|
||||
it('should not convert URLs when getFullFileUrl returns null', async () => {
|
||||
mockGetKeyFromFullUrl.mockResolvedValue('files/image-key.jpg');
|
||||
mockGetFullFileUrl.mockResolvedValue(null);
|
||||
|
||||
const ctx = createMockCtx();
|
||||
const input = createDefaultInput({
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
imageUrl: 'http://localhost:3000/f/file-id',
|
||||
},
|
||||
});
|
||||
|
||||
const caller = imageRouter.createCaller(ctx);
|
||||
const result = await caller.createImage(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(mockGetFullFileUrl).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -23,32 +23,9 @@ import {
|
|||
} from '@/types/asyncTask';
|
||||
import { generateUniqueSeeds } from '@/utils/number';
|
||||
|
||||
const log = debug('lobe-image:lambda');
|
||||
import { validateNoUrlsInConfig } from './utils';
|
||||
|
||||
/**
|
||||
* Recursively validate that no full URLs are present in the config
|
||||
* This is a defensive check to ensure only keys are stored in database
|
||||
*/
|
||||
function validateNoUrlsInConfig(obj: any, path: string = ''): void {
|
||||
if (typeof obj === 'string') {
|
||||
if (obj.startsWith('http://') || obj.startsWith('https://')) {
|
||||
throw new Error(
|
||||
`Invalid configuration: Found full URL instead of key at ${path || 'root'}. ` +
|
||||
`URL: "${obj.slice(0, 100)}${obj.length > 100 ? '...' : ''}". ` +
|
||||
`All URLs must be converted to storage keys before database insertion.`,
|
||||
);
|
||||
}
|
||||
} else if (Array.isArray(obj)) {
|
||||
obj.forEach((item, index) => {
|
||||
validateNoUrlsInConfig(item, `${path}[${index}]`);
|
||||
});
|
||||
} else if (obj && typeof obj === 'object') {
|
||||
Object.entries(obj).forEach(([key, value]) => {
|
||||
const currentPath = path ? `${path}.${key}` : key;
|
||||
validateNoUrlsInConfig(value, currentPath);
|
||||
});
|
||||
}
|
||||
}
|
||||
const log = debug('lobe-image:lambda');
|
||||
|
||||
const imageProcedure = authedProcedure
|
||||
.use(keyVaults)
|
||||
|
|
@ -103,11 +80,18 @@ export const imageRouter = router({
|
|||
if (Array.isArray(params.imageUrls) && params.imageUrls.length > 0) {
|
||||
log('Converting imageUrls to S3 keys for database storage: %O', params.imageUrls);
|
||||
try {
|
||||
const imageKeys = params.imageUrls.map((url) => {
|
||||
const key = fileService.getKeyFromFullUrl(url);
|
||||
log('Converted URL %s to key %s', url, key);
|
||||
return key;
|
||||
});
|
||||
const imageKeysWithNull = await Promise.all(
|
||||
params.imageUrls.map(async (url) => {
|
||||
const key = await fileService.getKeyFromFullUrl(url);
|
||||
if (key) {
|
||||
log('Converted URL %s to key %s', url, key);
|
||||
} else {
|
||||
log('Failed to extract key from URL: %s', url);
|
||||
}
|
||||
return key;
|
||||
}),
|
||||
);
|
||||
const imageKeys = imageKeysWithNull.filter((key): key is string => key !== null);
|
||||
|
||||
configForDatabase = {
|
||||
...configForDatabase,
|
||||
|
|
@ -115,29 +99,61 @@ export const imageRouter = router({
|
|||
};
|
||||
log('Successfully converted imageUrls to keys for database: %O', imageKeys);
|
||||
} catch (error) {
|
||||
log('Error converting imageUrls to keys: %O', error);
|
||||
log('Keeping original imageUrls due to conversion error');
|
||||
console.error('Error converting imageUrls to keys: %O', error);
|
||||
console.error('Keeping original imageUrls due to conversion error');
|
||||
}
|
||||
}
|
||||
// 2) Process single image in imageUrl
|
||||
if (typeof params.imageUrl === 'string' && params.imageUrl) {
|
||||
try {
|
||||
const key = fileService.getKeyFromFullUrl(params.imageUrl);
|
||||
log('Converted single imageUrl to key: %s -> %s', params.imageUrl, key);
|
||||
configForDatabase = { ...configForDatabase, imageUrl: key };
|
||||
const key = await fileService.getKeyFromFullUrl(params.imageUrl);
|
||||
if (key) {
|
||||
log('Converted single imageUrl to key: %s -> %s', params.imageUrl, key);
|
||||
configForDatabase = { ...configForDatabase, imageUrl: key };
|
||||
} else {
|
||||
log('Failed to extract key from single imageUrl: %s', params.imageUrl);
|
||||
}
|
||||
} catch (error) {
|
||||
log('Error converting imageUrl to key: %O', error);
|
||||
console.error('Error converting imageUrl to key: %O', error);
|
||||
// Keep original value if conversion fails
|
||||
}
|
||||
}
|
||||
|
||||
// In development, convert localhost proxy URLs to S3 URLs for async task access
|
||||
let generationParams = params;
|
||||
if (process.env.NODE_ENV === 'development') {
|
||||
const updates: Record<string, unknown> = {};
|
||||
|
||||
// Handle single imageUrl: localhost/f/{id} -> S3 URL
|
||||
if (typeof params.imageUrl === 'string' && params.imageUrl) {
|
||||
const s3Url = await fileService.getFullFileUrl(configForDatabase.imageUrl as string);
|
||||
if (s3Url) {
|
||||
log('Dev: converted proxy URL to S3 URL: %s -> %s', params.imageUrl, s3Url);
|
||||
updates.imageUrl = s3Url;
|
||||
}
|
||||
}
|
||||
|
||||
// Handle multiple imageUrls
|
||||
if (Array.isArray(params.imageUrls) && params.imageUrls.length > 0) {
|
||||
const s3Urls = await Promise.all(
|
||||
(configForDatabase.imageUrls as string[]).map((key) => fileService.getFullFileUrl(key)),
|
||||
);
|
||||
log('Dev: converted proxy URLs to S3 URLs: %O', s3Urls);
|
||||
updates.imageUrls = s3Urls;
|
||||
}
|
||||
|
||||
if (Object.keys(updates).length > 0) {
|
||||
generationParams = { ...params, ...updates };
|
||||
}
|
||||
}
|
||||
|
||||
// Defensive check: ensure no full URLs enter the database
|
||||
validateNoUrlsInConfig(configForDatabase, 'configForDatabase');
|
||||
|
||||
const chargeResult = await chargeBeforeGenerate({
|
||||
clientIp: ctx.clientIp,
|
||||
configForDatabase,
|
||||
generationParams: params,
|
||||
generationParams,
|
||||
generationTopicId,
|
||||
imageNum,
|
||||
model,
|
||||
|
|
@ -167,7 +183,7 @@ export const imageRouter = router({
|
|||
const [batch] = await tx.insert(generationBatches).values(newBatch).returning();
|
||||
log('Generation batch created successfully: %s', batch.id);
|
||||
|
||||
// 2. Create 4 generations (fixed at 4 images for phase one)
|
||||
// 2. Create generations
|
||||
const seeds =
|
||||
'seed' in params
|
||||
? generateUniqueSeeds(imageNum)
|
||||
|
|
@ -248,7 +264,7 @@ export const imageRouter = router({
|
|||
generationId: generation.id,
|
||||
generationTopicId,
|
||||
model,
|
||||
params,
|
||||
params: generationParams,
|
||||
provider,
|
||||
taskId: asyncTaskId,
|
||||
});
|
||||
|
|
@ -256,8 +272,8 @@ export const imageRouter = router({
|
|||
|
||||
log('All %d background async image generation tasks started', generationsWithTasks.length);
|
||||
} catch (e) {
|
||||
console.error('[createImage] Failed to process async tasks:', e);
|
||||
log('Failed to process async tasks: %O', e);
|
||||
console.error('Failed to process async tasks:', e);
|
||||
console.error('Failed to process async tasks: %O', e);
|
||||
|
||||
// If overall failure occurs, update all task statuses to failed
|
||||
try {
|
||||
|
|
@ -1,26 +1,6 @@
|
|||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
// Copy of the validation function from image.ts for testing
|
||||
function validateNoUrlsInConfig(obj: any, path: string = ''): void {
|
||||
if (typeof obj === 'string') {
|
||||
if (obj.startsWith('http://') || obj.startsWith('https://')) {
|
||||
throw new Error(
|
||||
`Invalid configuration: Found full URL instead of key at ${path || 'root'}. ` +
|
||||
`URL: "${obj.slice(0, 100)}${obj.length > 100 ? '...' : ''}". ` +
|
||||
`All URLs must be converted to storage keys before database insertion.`,
|
||||
);
|
||||
}
|
||||
} else if (Array.isArray(obj)) {
|
||||
obj.forEach((item, index) => {
|
||||
validateNoUrlsInConfig(item, `${path}[${index}]`);
|
||||
});
|
||||
} else if (obj && typeof obj === 'object') {
|
||||
Object.entries(obj).forEach(([key, value]) => {
|
||||
const currentPath = path ? `${path}.${key}` : key;
|
||||
validateNoUrlsInConfig(value, currentPath);
|
||||
});
|
||||
}
|
||||
}
|
||||
import { validateNoUrlsInConfig } from './utils';
|
||||
|
||||
describe('imageRouter', () => {
|
||||
describe('validateNoUrlsInConfig utility', () => {
|
||||
24
src/server/routers/lambda/image/utils.ts
Normal file
24
src/server/routers/lambda/image/utils.ts
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
/**
|
||||
* Recursively validate that no full URLs are present in the config
|
||||
* This is a defensive check to ensure only keys are stored in database
|
||||
*/
|
||||
export function validateNoUrlsInConfig(obj: any, path: string = ''): void {
|
||||
if (typeof obj === 'string') {
|
||||
if (obj.startsWith('http://') || obj.startsWith('https://')) {
|
||||
throw new Error(
|
||||
`Invalid configuration: Found full URL instead of key at ${path || 'root'}. ` +
|
||||
`URL: "${obj.slice(0, 100)}${obj.length > 100 ? '...' : ''}". ` +
|
||||
`All URLs must be converted to storage keys before database insertion.`,
|
||||
);
|
||||
}
|
||||
} else if (Array.isArray(obj)) {
|
||||
obj.forEach((item, index) => {
|
||||
validateNoUrlsInConfig(item, `${path}[${index}]`);
|
||||
});
|
||||
} else if (obj && typeof obj === 'object') {
|
||||
Object.entries(obj).forEach(([key, value]) => {
|
||||
const currentPath = path ? `${path}.${key}` : key;
|
||||
validateNoUrlsInConfig(value, currentPath);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -231,12 +231,12 @@ describe('FileService', () => {
|
|||
expect(result).toBe(expectedUrl);
|
||||
});
|
||||
|
||||
it('should delegate getKeyFromFullUrl to implementation', () => {
|
||||
it('should delegate getKeyFromFullUrl to implementation', async () => {
|
||||
const testUrl = 'https://example.com/path/to/file.jpg';
|
||||
const expectedKey = 'path/to/file.jpg';
|
||||
vi.mocked(service['impl'].getKeyFromFullUrl).mockReturnValue(expectedKey);
|
||||
vi.mocked(service['impl'].getKeyFromFullUrl).mockResolvedValue(expectedKey);
|
||||
|
||||
const result = service.getKeyFromFullUrl(testUrl);
|
||||
const result = await service.getKeyFromFullUrl(testUrl);
|
||||
|
||||
expect(service['impl'].getKeyFromFullUrl).toHaveBeenCalledWith(testUrl);
|
||||
expect(result).toBe(expectedKey);
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import { LobeChatDatabase } from '@lobechat/database';
|
||||
|
||||
import { S3StaticFileImpl } from './s3';
|
||||
import { type FileServiceImpl } from './type';
|
||||
|
||||
|
|
@ -5,8 +7,6 @@ import { type FileServiceImpl } from './type';
|
|||
* Create file service module
|
||||
* Returns S3 file implementation for cloud storage
|
||||
*/
|
||||
export const createFileServiceModule = (): FileServiceImpl => {
|
||||
return new S3StaticFileImpl();
|
||||
export const createFileServiceModule = (db: LobeChatDatabase): FileServiceImpl => {
|
||||
return new S3StaticFileImpl(db);
|
||||
};
|
||||
|
||||
export type { FileServiceImpl } from './type';
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { FileModel } from '@/database/models/file';
|
||||
|
||||
import { S3StaticFileImpl } from './s3';
|
||||
|
||||
const config = {
|
||||
|
|
@ -33,11 +35,14 @@ vi.mock('@/server/modules/S3', () => ({
|
|||
})),
|
||||
}));
|
||||
|
||||
// Mock db
|
||||
const mockDb = {} as any;
|
||||
|
||||
describe('S3StaticFileImpl', () => {
|
||||
let fileService: S3StaticFileImpl;
|
||||
|
||||
beforeEach(() => {
|
||||
fileService = new S3StaticFileImpl();
|
||||
fileService = new S3StaticFileImpl(mockDb);
|
||||
});
|
||||
|
||||
describe('getFullFileUrl', () => {
|
||||
|
|
@ -74,7 +79,7 @@ describe('S3StaticFileImpl', () => {
|
|||
const fullUrl = 'https://s3.example.com/bucket/path/to/file.jpg?X-Amz-Signature=expired';
|
||||
|
||||
// Mock getKeyFromFullUrl to return the extracted key
|
||||
vi.spyOn(fileService, 'getKeyFromFullUrl').mockReturnValue('path/to/file.jpg');
|
||||
vi.spyOn(fileService, 'getKeyFromFullUrl').mockResolvedValue('path/to/file.jpg');
|
||||
|
||||
const result = await fileService.getFullFileUrl(fullUrl);
|
||||
|
||||
|
|
@ -86,7 +91,7 @@ describe('S3StaticFileImpl', () => {
|
|||
it('should handle full URL input by extracting key (S3_SET_ACL=true)', async () => {
|
||||
const fullUrl = 'https://s3.example.com/bucket/path/to/file.jpg';
|
||||
|
||||
vi.spyOn(fileService, 'getKeyFromFullUrl').mockReturnValue('path/to/file.jpg');
|
||||
vi.spyOn(fileService, 'getKeyFromFullUrl').mockResolvedValue('path/to/file.jpg');
|
||||
|
||||
const result = await fileService.getFullFileUrl(fullUrl);
|
||||
|
||||
|
|
@ -108,13 +113,23 @@ describe('S3StaticFileImpl', () => {
|
|||
it('should handle http:// URLs for legacy compatibility', async () => {
|
||||
const httpUrl = 'http://s3.example.com/bucket/path/to/file.jpg';
|
||||
|
||||
vi.spyOn(fileService, 'getKeyFromFullUrl').mockReturnValue('path/to/file.jpg');
|
||||
vi.spyOn(fileService, 'getKeyFromFullUrl').mockResolvedValue('path/to/file.jpg');
|
||||
|
||||
const result = await fileService.getFullFileUrl(httpUrl);
|
||||
|
||||
expect(fileService.getKeyFromFullUrl).toHaveBeenCalledWith(httpUrl);
|
||||
expect(result).toBe('https://example.com/path/to/file.jpg');
|
||||
});
|
||||
|
||||
it('should throw error when key extraction returns null', async () => {
|
||||
const fullUrl = 'https://s3.example.com/f/nonexistent';
|
||||
|
||||
vi.spyOn(fileService, 'getKeyFromFullUrl').mockResolvedValue(null);
|
||||
|
||||
await expect(fileService.getFullFileUrl(fullUrl)).rejects.toThrow(
|
||||
'Key not found from url: ' + fullUrl,
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -179,63 +194,66 @@ describe('S3StaticFileImpl', () => {
|
|||
});
|
||||
|
||||
describe('getKeyFromFullUrl', () => {
|
||||
it('当S3_ENABLE_PATH_STYLE为false时应该正确提取key', () => {
|
||||
it('should extract fileId from proxy URL and return S3 key from database', async () => {
|
||||
const proxyUrl = 'http://localhost:3010/f/abc123';
|
||||
const expectedKey = 'ppp/491067/image.jpg';
|
||||
|
||||
vi.spyOn(FileModel, 'getFileById').mockResolvedValue({ url: expectedKey } as any);
|
||||
|
||||
const result = await fileService.getKeyFromFullUrl(proxyUrl);
|
||||
|
||||
expect(FileModel.getFileById).toHaveBeenCalledWith(mockDb, 'abc123');
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
|
||||
it('should return null when file is not found in database', async () => {
|
||||
const proxyUrl = 'http://localhost:3010/f/nonexistent';
|
||||
|
||||
vi.spyOn(FileModel, 'getFileById').mockResolvedValue(undefined);
|
||||
|
||||
const result = await fileService.getKeyFromFullUrl(proxyUrl);
|
||||
|
||||
expect(FileModel.getFileById).toHaveBeenCalledWith(mockDb, 'nonexistent');
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should handle URL with different domain', async () => {
|
||||
const proxyUrl = 'https://example.com/f/file456';
|
||||
const expectedKey = 'uploads/file.png';
|
||||
|
||||
vi.spyOn(FileModel, 'getFileById').mockResolvedValue({ url: expectedKey } as any);
|
||||
|
||||
const result = await fileService.getKeyFromFullUrl(proxyUrl);
|
||||
|
||||
expect(FileModel.getFileById).toHaveBeenCalledWith(mockDb, 'file456');
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
|
||||
it('should extract key from legacy S3 URL (non /f/ path)', async () => {
|
||||
const s3Url = 'https://example.com/path/to/file.jpg';
|
||||
|
||||
const result = await fileService.getKeyFromFullUrl(s3Url);
|
||||
|
||||
// Legacy S3 URL: extract key from pathname
|
||||
expect(result).toBe('path/to/file.jpg');
|
||||
});
|
||||
|
||||
it('should extract key with path-style S3 URL', async () => {
|
||||
config.S3_ENABLE_PATH_STYLE = true;
|
||||
const s3Url = 'https://example.com/my-bucket/path/to/file.jpg';
|
||||
|
||||
const result = await fileService.getKeyFromFullUrl(s3Url);
|
||||
|
||||
expect(result).toBe('path/to/file.jpg');
|
||||
config.S3_ENABLE_PATH_STYLE = false;
|
||||
const url = 'https://example.com/path/to/file.jpg';
|
||||
|
||||
const result = fileService.getKeyFromFullUrl(url);
|
||||
|
||||
expect(result).toBe('path/to/file.jpg');
|
||||
config.S3_ENABLE_PATH_STYLE = false; // reset
|
||||
});
|
||||
|
||||
it('当S3_ENABLE_PATH_STYLE为true时应该正确提取key', () => {
|
||||
config.S3_ENABLE_PATH_STYLE = true;
|
||||
const url = 'https://example.com/my-bucket/path/to/file.jpg';
|
||||
|
||||
const result = fileService.getKeyFromFullUrl(url);
|
||||
|
||||
expect(result).toBe('path/to/file.jpg');
|
||||
config.S3_ENABLE_PATH_STYLE = false; // reset
|
||||
});
|
||||
|
||||
it('当S3_ENABLE_PATH_STYLE为true但缺少bucket名称时应该返回pathname', () => {
|
||||
config.S3_ENABLE_PATH_STYLE = true;
|
||||
config.S3_BUCKET = '';
|
||||
const url = 'https://example.com/path/to/file.jpg';
|
||||
|
||||
const result = fileService.getKeyFromFullUrl(url);
|
||||
|
||||
expect(result).toBe('path/to/file.jpg');
|
||||
config.S3_ENABLE_PATH_STYLE = false; // reset
|
||||
config.S3_BUCKET = 'my-bucket'; // reset
|
||||
});
|
||||
|
||||
it('当URL格式不正确时应该返回原始字符串', () => {
|
||||
it('should return null for invalid URL', async () => {
|
||||
const invalidUrl = 'not-a-valid-url';
|
||||
|
||||
const result = fileService.getKeyFromFullUrl(invalidUrl);
|
||||
const result = await fileService.getKeyFromFullUrl(invalidUrl);
|
||||
|
||||
expect(result).toBe('not-a-valid-url');
|
||||
});
|
||||
|
||||
it('应该处理根路径文件', () => {
|
||||
config.S3_ENABLE_PATH_STYLE = false;
|
||||
const url = 'https://example.com/file.jpg';
|
||||
|
||||
const result = fileService.getKeyFromFullUrl(url);
|
||||
|
||||
expect(result).toBe('file.jpg');
|
||||
});
|
||||
|
||||
it('当path-style URL路径格式不符合预期时应该使用fallback', () => {
|
||||
config.S3_ENABLE_PATH_STYLE = true;
|
||||
const url = 'https://example.com/unexpected/path/file.jpg';
|
||||
|
||||
const result = fileService.getKeyFromFullUrl(url);
|
||||
|
||||
expect(result).toBe('unexpected/path/file.jpg');
|
||||
config.S3_ENABLE_PATH_STYLE = false; // reset
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -1,18 +1,21 @@
|
|||
import { LobeChatDatabase } from '@lobechat/database';
|
||||
import urlJoin from 'url-join';
|
||||
|
||||
import { FileModel } from '@/database/models/file';
|
||||
import { fileEnv } from '@/envs/file';
|
||||
import { FileS3 } from '@/server/modules/S3';
|
||||
|
||||
import { type FileServiceImpl } from './type';
|
||||
import { extractKeyFromUrlOrReturnOriginal } from './utils';
|
||||
|
||||
/**
|
||||
* S3-based file service implementation
|
||||
*/
|
||||
export class S3StaticFileImpl implements FileServiceImpl {
|
||||
private readonly s3: FileS3;
|
||||
private readonly db: LobeChatDatabase;
|
||||
|
||||
constructor() {
|
||||
constructor(db: LobeChatDatabase) {
|
||||
this.db = db;
|
||||
this.s3 = new FileS3();
|
||||
}
|
||||
|
||||
|
|
@ -51,8 +54,16 @@ export class S3StaticFileImpl implements FileServiceImpl {
|
|||
async getFullFileUrl(url?: string | null, expiresIn?: number): Promise<string> {
|
||||
if (!url) return '';
|
||||
|
||||
// Handle legacy data compatibility using shared utility
|
||||
const key = extractKeyFromUrlOrReturnOriginal(url, this.getKeyFromFullUrl.bind(this));
|
||||
// Handle legacy data compatibility - extract key from full URL if needed
|
||||
// Related issue: https://github.com/lobehub/lobe-chat/issues/8994
|
||||
let key = url;
|
||||
if (url.startsWith('http://') || url.startsWith('https://')) {
|
||||
const extractedKey = await this.getKeyFromFullUrl(url);
|
||||
if (!extractedKey) {
|
||||
throw new Error('Key not found from url: ' + url);
|
||||
}
|
||||
key = extractedKey;
|
||||
}
|
||||
|
||||
// If bucket is not set public read, the preview address needs to be regenerated each time
|
||||
if (!fileEnv.S3_SET_ACL) {
|
||||
|
|
@ -66,38 +77,35 @@ export class S3StaticFileImpl implements FileServiceImpl {
|
|||
return urlJoin(fileEnv.S3_PUBLIC_DOMAIN!, key);
|
||||
}
|
||||
|
||||
getKeyFromFullUrl(url: string): string {
|
||||
async getKeyFromFullUrl(url: string): Promise<string | null> {
|
||||
try {
|
||||
const urlObject = new URL(url);
|
||||
const { pathname } = urlObject;
|
||||
|
||||
let key: string;
|
||||
|
||||
if (fileEnv.S3_ENABLE_PATH_STYLE) {
|
||||
if (!fileEnv.S3_BUCKET) {
|
||||
// In path-style, we need bucket name to extract key
|
||||
// but if not provided, we can only guess the key is the pathname
|
||||
return pathname.startsWith('/') ? pathname.slice(1) : pathname;
|
||||
}
|
||||
// For path-style URLs, the path is /<bucket>/<key>
|
||||
// We need to remove the leading slash and the bucket name.
|
||||
const bucketPrefix = `/${fileEnv.S3_BUCKET}/`;
|
||||
if (pathname.startsWith(bucketPrefix)) {
|
||||
key = pathname.slice(bucketPrefix.length);
|
||||
} else {
|
||||
// Fallback for unexpected path format
|
||||
key = pathname.startsWith('/') ? pathname.slice(1) : pathname;
|
||||
}
|
||||
} else {
|
||||
// For virtual-hosted-style URLs, the path is /<key>
|
||||
// We just need to remove the leading slash.
|
||||
key = pathname.slice(1);
|
||||
// Case 1: File proxy URL pattern /f/{fileId} - query database for S3 key
|
||||
if (pathname.startsWith('/f/')) {
|
||||
const fileId = pathname.slice(3); // Remove '/f/' prefix
|
||||
const file = await FileModel.getFileById(this.db, fileId);
|
||||
return file?.url ?? null;
|
||||
}
|
||||
|
||||
return key;
|
||||
// Case 2: Legacy S3 URL - extract key from pathname
|
||||
if (fileEnv.S3_ENABLE_PATH_STYLE) {
|
||||
if (!fileEnv.S3_BUCKET) {
|
||||
return pathname.startsWith('/') ? pathname.slice(1) : pathname;
|
||||
}
|
||||
const bucketPrefix = `/${fileEnv.S3_BUCKET}/`;
|
||||
if (pathname.startsWith(bucketPrefix)) {
|
||||
return pathname.slice(bucketPrefix.length);
|
||||
}
|
||||
return pathname.startsWith('/') ? pathname.slice(1) : pathname;
|
||||
}
|
||||
|
||||
// Virtual-hosted-style: path is /<key>
|
||||
return pathname.slice(1);
|
||||
} catch {
|
||||
// if url is not a valid URL, it may be a key itself
|
||||
return url;
|
||||
// If url is not a valid URL, return null
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@ export interface FileServiceImpl {
|
|||
* Create pre-signed upload URL
|
||||
*/
|
||||
createPreSignedUrl(key: string): Promise<string>;
|
||||
|
||||
/**
|
||||
* Create pre-signed preview URL
|
||||
*/
|
||||
|
|
@ -46,7 +45,7 @@ export interface FileServiceImpl {
|
|||
/**
|
||||
* Extract key from full URL
|
||||
*/
|
||||
getKeyFromFullUrl(url: string): string;
|
||||
getKeyFromFullUrl(url: string): Promise<string | null>;
|
||||
|
||||
/**
|
||||
* Upload content
|
||||
|
|
|
|||
|
|
@ -1,154 +0,0 @@
|
|||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { extractKeyFromUrlOrReturnOriginal } from './utils';
|
||||
|
||||
describe('extractKeyFromUrlOrReturnOriginal', () => {
|
||||
const mockGetKeyFromFullUrl = vi.fn();
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('URL detection', () => {
|
||||
it('should detect https:// URLs and extract key', () => {
|
||||
const httpsUrl = 'https://s3.example.com/bucket/path/to/file.jpg';
|
||||
const expectedKey = 'path/to/file.jpg';
|
||||
|
||||
mockGetKeyFromFullUrl.mockReturnValue(expectedKey);
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(httpsUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(httpsUrl);
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
|
||||
it('should detect http:// URLs and extract key', () => {
|
||||
const httpUrl = 'http://s3.example.com/bucket/path/to/file.jpg';
|
||||
const expectedKey = 'path/to/file.jpg';
|
||||
|
||||
mockGetKeyFromFullUrl.mockReturnValue(expectedKey);
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(httpUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(httpUrl);
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
|
||||
it('should handle presigned URLs with query parameters', () => {
|
||||
const presignedUrl = 'https://s3.amazonaws.com/bucket/file.jpg?X-Amz-Signature=abc&X-Amz-Expires=3600';
|
||||
const expectedKey = 'file.jpg';
|
||||
|
||||
mockGetKeyFromFullUrl.mockReturnValue(expectedKey);
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(presignedUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(presignedUrl);
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
});
|
||||
|
||||
describe('key passthrough', () => {
|
||||
it('should return original string for plain keys', () => {
|
||||
const key = 'path/to/file.jpg';
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(key, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).not.toHaveBeenCalled();
|
||||
expect(result).toBe(key);
|
||||
});
|
||||
|
||||
it('should return original string for desktop:// keys', () => {
|
||||
const desktopKey = 'desktop://documents/file.pdf';
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(desktopKey, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).not.toHaveBeenCalled();
|
||||
expect(result).toBe(desktopKey);
|
||||
});
|
||||
|
||||
it('should return original string for relative paths', () => {
|
||||
const relativePath = './assets/image.png';
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(relativePath, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).not.toHaveBeenCalled();
|
||||
expect(result).toBe(relativePath);
|
||||
});
|
||||
|
||||
it('should return original string for file:// protocol', () => {
|
||||
const fileUrl = 'file:///Users/test/file.txt';
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(fileUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).not.toHaveBeenCalled();
|
||||
expect(result).toBe(fileUrl);
|
||||
});
|
||||
});
|
||||
|
||||
describe('edge cases', () => {
|
||||
it('should handle empty string', () => {
|
||||
const emptyString = '';
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(emptyString, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).not.toHaveBeenCalled();
|
||||
expect(result).toBe(emptyString);
|
||||
});
|
||||
|
||||
it('should handle strings that start with http but are not URLs', () => {
|
||||
const notUrl = 'httpish-string';
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(notUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).not.toHaveBeenCalled();
|
||||
expect(result).toBe(notUrl);
|
||||
});
|
||||
|
||||
it('should handle case sensitivity correctly', () => {
|
||||
const upperCaseHttps = 'HTTPS://example.com/file.jpg';
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(upperCaseHttps, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).not.toHaveBeenCalled();
|
||||
expect(result).toBe(upperCaseHttps);
|
||||
});
|
||||
});
|
||||
|
||||
describe('legacy bug scenarios', () => {
|
||||
it('should handle S3 public URLs', () => {
|
||||
const s3PublicUrl = 'https://mybucket.s3.amazonaws.com/images/photo.jpg';
|
||||
const expectedKey = 'images/photo.jpg';
|
||||
|
||||
mockGetKeyFromFullUrl.mockReturnValue(expectedKey);
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(s3PublicUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(s3PublicUrl);
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
|
||||
it('should handle custom domain S3 URLs', () => {
|
||||
const customDomainUrl = 'https://cdn.example.com/files/document.pdf';
|
||||
const expectedKey = 'files/document.pdf';
|
||||
|
||||
mockGetKeyFromFullUrl.mockReturnValue(expectedKey);
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(customDomainUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(customDomainUrl);
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
|
||||
it('should handle local development URLs', () => {
|
||||
const localUrl = 'http://localhost:3000/desktop-file/images/screenshot.png';
|
||||
const expectedKey = 'desktop://images/screenshot.png';
|
||||
|
||||
mockGetKeyFromFullUrl.mockReturnValue(expectedKey);
|
||||
|
||||
const result = extractKeyFromUrlOrReturnOriginal(localUrl, mockGetKeyFromFullUrl);
|
||||
|
||||
expect(mockGetKeyFromFullUrl).toHaveBeenCalledWith(localUrl);
|
||||
expect(result).toBe(expectedKey);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,17 +0,0 @@
|
|||
/**
|
||||
* Handle legacy bug where full URLs were stored instead of keys
|
||||
* Some historical data stored complete URLs in database instead of just keys
|
||||
* Related issue: https://github.com/lobehub/lobe-chat/issues/8994
|
||||
*/
|
||||
export function extractKeyFromUrlOrReturnOriginal(
|
||||
url: string,
|
||||
getKeyFromFullUrl: (url: string) => string,
|
||||
): string {
|
||||
// Only process URLs that start with http:// or https://
|
||||
if (url.startsWith('http://') || url.startsWith('https://')) {
|
||||
// Extract key from full URL for legacy data compatibility
|
||||
return getKeyFromFullUrl(url);
|
||||
}
|
||||
// Return original input if it's already a key
|
||||
return url;
|
||||
}
|
||||
|
|
@ -9,7 +9,8 @@ import { type FileItem } from '@/database/schemas';
|
|||
import { appEnv } from '@/envs/app';
|
||||
import { TempFileManager } from '@/server/utils/tempFileManager';
|
||||
|
||||
import { type FileServiceImpl, createFileServiceModule } from './impls';
|
||||
import { createFileServiceModule } from './impls';
|
||||
import { type FileServiceImpl } from './impls/type';
|
||||
|
||||
/**
|
||||
* File service class
|
||||
|
|
@ -19,11 +20,12 @@ export class FileService {
|
|||
private userId: string;
|
||||
private fileModel: FileModel;
|
||||
|
||||
private impl: FileServiceImpl = createFileServiceModule();
|
||||
private impl: FileServiceImpl;
|
||||
|
||||
constructor(db: LobeChatDatabase, userId: string) {
|
||||
this.userId = userId;
|
||||
this.fileModel = new FileModel(db, userId);
|
||||
this.impl = createFileServiceModule(db);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -95,7 +97,7 @@ export class FileService {
|
|||
/**
|
||||
* Extract key from full URL
|
||||
*/
|
||||
public getKeyFromFullUrl(url: string): string {
|
||||
public async getKeyFromFullUrl(url: string): Promise<string | null> {
|
||||
return this.impl.getKeyFromFullUrl(url);
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue