Skip to content

Commit 2b198b4

Browse files
arvinxxramu-narasinga
authored andcommitted
🐛 fix: fix inbox agent can not save config (lobehub#6186)
1 parent ea0cc34 commit 2b198b4

File tree

11 files changed

+174
-49
lines changed

11 files changed

+174
-49
lines changed

src/database/server/models/__tests__/session.test.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -628,7 +628,7 @@ describe('SessionModel', () => {
628628

629629
describe('createInbox', () => {
630630
it('should create inbox session if not exists', async () => {
631-
const inbox = await sessionModel.createInbox();
631+
const inbox = await sessionModel.createInbox({});
632632

633633
expect(inbox).toBeDefined();
634634
expect(inbox?.slug).toBe('inbox');
@@ -641,10 +641,10 @@ describe('SessionModel', () => {
641641

642642
it('should not create duplicate inbox session', async () => {
643643
// Create first inbox
644-
await sessionModel.createInbox();
644+
await sessionModel.createInbox({});
645645

646646
// Try to create another inbox
647-
const duplicateInbox = await sessionModel.createInbox();
647+
const duplicateInbox = await sessionModel.createInbox({});
648648

649649
// Should return undefined as inbox already exists
650650
expect(duplicateInbox).toBeUndefined();

src/database/server/models/session.ts

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { Column, count, sql } from 'drizzle-orm';
22
import { and, asc, desc, eq, gt, inArray, isNull, like, not, or } from 'drizzle-orm/expressions';
3+
import { DeepPartial } from 'utility-types';
34

4-
import { appEnv } from '@/config/app';
55
import { DEFAULT_INBOX_AVATAR } from '@/const/meta';
66
import { INBOX_SESSION_ID } from '@/const/session';
77
import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
@@ -13,7 +13,7 @@ import {
1313
genWhere,
1414
} from '@/database/utils/genWhere';
1515
import { idGenerator } from '@/database/utils/idGenerator';
16-
import { parseAgentConfig } from '@/server/globalConfig/parseDefaultAgent';
16+
import { LobeAgentConfig } from '@/types/agent';
1717
import { ChatSessionList, LobeAgentSession, SessionRankItem } from '@/types/session';
1818
import { merge } from '@/utils/merge';
1919

@@ -226,16 +226,15 @@ export class SessionModel {
226226
});
227227
};
228228

229-
createInbox = async () => {
229+
createInbox = async (defaultAgentConfig: DeepPartial<LobeAgentConfig>) => {
230230
const item = await this.db.query.sessions.findFirst({
231231
where: and(eq(sessions.userId, this.userId), eq(sessions.slug, INBOX_SESSION_ID)),
232232
});
233-
if (item) return;
234233

235-
const serverAgentConfig = parseAgentConfig(appEnv.DEFAULT_AGENT_CONFIG) || {};
234+
if (item) return;
236235

237236
return await this.create({
238-
config: merge(DEFAULT_AGENT_CONFIG, serverAgentConfig),
237+
config: merge(DEFAULT_AGENT_CONFIG, defaultAgentConfig),
239238
slug: INBOX_SESSION_ID,
240239
type: 'agent',
241240
});

src/database/server/models/user.ts

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import { merge } from '@/utils/merge';
1010
import { today } from '@/utils/time';
1111

1212
import { NewUser, UserItem, UserSettingsItem, userSettings, users } from '../../schemas';
13-
import { SessionModel } from './session';
1413

1514
type DecryptUserKeyVaults = (
1615
encryptKeyVaultsStr: string | null,
@@ -160,10 +159,7 @@ export class UserModel {
160159
.values({ ...params })
161160
.returning();
162161

163-
// Create an inbox session for the user
164-
const model = new SessionModel(db, user.id);
165-
166-
await model.createInbox();
162+
return user;
167163
};
168164

169165
static deleteUser = async (db: LobeChatDatabase, id: string) => {

src/features/DevPanel/features/Table/TableCell.tsx

Lines changed: 34 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,27 @@
1-
import { Typography } from 'antd';
2-
import { createStyles } from 'antd-style';
31
import dayjs from 'dayjs';
42
import { get, isDate } from 'lodash-es';
53
import React, { useMemo } from 'react';
64

7-
import TooltipContent from './TooltipContent';
5+
// import TooltipContent from './TooltipContent';
86

9-
const { Text } = Typography;
7+
// const { Text } = Typography;
108

11-
const useStyles = createStyles(({ token, css }) => ({
12-
cell: css`
13-
font-family: ${token.fontFamilyCode};
14-
font-size: ${token.fontSizeSM}px;
15-
`,
16-
tooltip: css`
17-
border: 1px solid ${token.colorBorder};
18-
19-
font-family: ${token.fontFamilyCode};
20-
font-size: ${token.fontSizeSM}px;
21-
color: ${token.colorText} !important;
22-
word-break: break-all;
23-
24-
background: ${token.colorBgElevated} !important;
25-
`,
26-
}));
9+
// const useStyles = createStyles(({ token, css }) => ({
10+
// cell: css`
11+
// font-family: ${token.fontFamilyCode};
12+
// font-size: ${token.fontSizeSM}px;
13+
// `,
14+
// tooltip: css`
15+
// border: 1px solid ${token.colorBorder};
16+
//
17+
// font-family: ${token.fontFamilyCode};
18+
// font-size: ${token.fontSizeSM}px;
19+
// color: ${token.colorText} !important;
20+
// word-break: break-all;
21+
//
22+
// background: ${token.colorBgElevated} !important;
23+
// `,
24+
// }));
2725

2826
interface TableCellProps {
2927
column: string;
@@ -32,7 +30,7 @@ interface TableCellProps {
3230
}
3331

3432
const TableCell = ({ dataItem, column, rowIndex }: TableCellProps) => {
35-
const { styles } = useStyles();
33+
// const { styles } = useStyles();
3634
const data = get(dataItem, column);
3735
const content = useMemo(() => {
3836
if (isDate(data)) return dayjs(data).format('YYYY-MM-DD HH:mm:ss');
@@ -54,18 +52,21 @@ const TableCell = ({ dataItem, column, rowIndex }: TableCellProps) => {
5452

5553
return (
5654
<td key={column} onDoubleClick={() => console.log('Edit cell:', rowIndex, column)}>
57-
<Text
58-
className={styles.cell}
59-
ellipsis={{
60-
tooltip: {
61-
arrow: false,
62-
classNames: { body: styles.tooltip },
63-
title: <TooltipContent>{content}</TooltipContent>,
64-
},
65-
}}
66-
>
67-
{content}
68-
</Text>
55+
{content}
56+
57+
{/* 不能使用 antd 的 Text, 会有大量的重渲染导致滚动极其卡顿 */}
58+
{/*<Text*/}
59+
{/* className={styles.cell}*/}
60+
{/* ellipsis={{*/}
61+
{/* tooltip: {*/}
62+
{/* arrow: false,*/}
63+
{/* classNames: { body: styles.tooltip },*/}
64+
{/* title: <TooltipContent>{content}</TooltipContent>,*/}
65+
{/* },*/}
66+
{/* }}*/}
67+
{/*>*/}
68+
{/* {content}*/}
69+
{/*</Text>*/}
6970
</td>
7071
);
7172
};

src/libs/next-auth/adapter/index.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { Adapter, AdapterAccount } from 'next-auth/adapters';
1010

1111
import * as schema from '@/database/schemas';
1212
import { UserModel } from '@/database/server/models/user';
13+
import { AgentService } from '@/server/services/agent';
1314
import { merge } from '@/utils/merge';
1415

1516
import {
@@ -65,6 +66,7 @@ export function LobeNextAuthDbAdapter(serverDB: NeonDatabase<typeof schema>): Ad
6566
const adapterUser = mapLobeUserToAdapterUser(existingUser);
6667
return adapterUser;
6768
}
69+
6870
// create a new user if it does not exist
6971
await UserModel.createUser(
7072
serverDB,
@@ -77,6 +79,11 @@ export function LobeNextAuthDbAdapter(serverDB: NeonDatabase<typeof schema>): Ad
7779
name,
7880
}),
7981
);
82+
83+
// 3. Create an inbox session for the user
84+
const agentService = new AgentService(serverDB, id);
85+
await agentService.createInbox();
86+
8087
return { ...user, id: providerAccountId ?? id };
8188
},
8289
async createVerificationToken(data): Promise<VerificationToken | null | undefined> {

src/server/routers/lambda/agent.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { SessionModel } from '@/database/server/models/session';
1010
import { UserModel } from '@/database/server/models/user';
1111
import { pino } from '@/libs/logger';
1212
import { authedProcedure, router } from '@/libs/trpc';
13+
import { AgentService } from '@/server/services/agent';
1314
import { KnowledgeItem, KnowledgeType } from '@/types/knowledgeBase';
1415

1516
const agentProcedure = authedProcedure.use(async (opts) => {
@@ -18,6 +19,7 @@ const agentProcedure = authedProcedure.use(async (opts) => {
1819
return opts.next({
1920
ctx: {
2021
agentModel: new AgentModel(serverDB, ctx.userId),
22+
agentService: new AgentService(serverDB, ctx.userId),
2123
fileModel: new FileModel(serverDB, ctx.userId),
2224
knowledgeBaseModel: new KnowledgeBaseModel(serverDB, ctx.userId),
2325
sessionModel: new SessionModel(serverDB, ctx.userId),
@@ -91,7 +93,7 @@ export const agentRouter = router({
9193
const user = await UserModel.findById(serverDB, ctx.userId);
9294
if (!user) return DEFAULT_AGENT_CONFIG;
9395

94-
const res = await ctx.sessionModel.createInbox();
96+
const res = await ctx.agentService.createInbox();
9597
pino.info('create inbox session', res);
9698
}
9799
}
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
// @vitest-environment node
2+
import { beforeEach, describe, expect, it, vi } from 'vitest';
3+
4+
import { SessionModel } from '@/database/server/models/session';
5+
import { parseAgentConfig } from '@/server/globalConfig/parseDefaultAgent';
6+
7+
import { AgentService } from './index';
8+
9+
vi.mock('@/config/app', () => ({
10+
appEnv: {
11+
DEFAULT_AGENT_CONFIG: 'model=gpt-4;temperature=0.7',
12+
},
13+
}));
14+
15+
vi.mock('@/server/globalConfig/parseDefaultAgent', () => ({
16+
parseAgentConfig: vi.fn(),
17+
}));
18+
19+
vi.mock('@/database/server/models/session', () => ({
20+
SessionModel: vi.fn(),
21+
}));
22+
23+
describe('AgentService', () => {
24+
let service: AgentService;
25+
const mockDb = {} as any;
26+
const mockUserId = 'test-user-id';
27+
28+
beforeEach(() => {
29+
vi.clearAllMocks();
30+
service = new AgentService(mockDb, mockUserId);
31+
});
32+
33+
describe('createInbox', () => {
34+
it('should create inbox with default agent config', async () => {
35+
const mockConfig = { model: 'gpt-4', temperature: 0.7 };
36+
const mockSessionModel = {
37+
createInbox: vi.fn(),
38+
};
39+
40+
(SessionModel as any).mockImplementation(() => mockSessionModel);
41+
(parseAgentConfig as any).mockReturnValue(mockConfig);
42+
43+
await service.createInbox();
44+
45+
expect(SessionModel).toHaveBeenCalledWith(mockDb, mockUserId);
46+
expect(parseAgentConfig).toHaveBeenCalledWith('model=gpt-4;temperature=0.7');
47+
expect(mockSessionModel.createInbox).toHaveBeenCalledWith(mockConfig);
48+
});
49+
50+
it('should create inbox with empty config if parseAgentConfig returns undefined', async () => {
51+
const mockSessionModel = {
52+
createInbox: vi.fn(),
53+
};
54+
55+
(SessionModel as any).mockImplementation(() => mockSessionModel);
56+
(parseAgentConfig as any).mockReturnValue(undefined);
57+
58+
await service.createInbox();
59+
60+
expect(SessionModel).toHaveBeenCalledWith(mockDb, mockUserId);
61+
expect(parseAgentConfig).toHaveBeenCalledWith('model=gpt-4;temperature=0.7');
62+
expect(mockSessionModel.createInbox).toHaveBeenCalledWith({});
63+
});
64+
});
65+
});

src/server/services/agent/index.ts

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import { appEnv } from '@/config/app';
2+
import { SessionModel } from '@/database/server/models/session';
3+
import { LobeChatDatabase } from '@/database/type';
4+
import { parseAgentConfig } from '@/server/globalConfig/parseDefaultAgent';
5+
6+
export class AgentService {
7+
private readonly userId: string;
8+
private readonly db: LobeChatDatabase;
9+
10+
constructor(db: LobeChatDatabase, userId: string) {
11+
this.userId = userId;
12+
this.db = db;
13+
}
14+
15+
async createInbox() {
16+
const sessionModel = new SessionModel(this.db, this.userId);
17+
18+
const defaultAgentConfig = parseAgentConfig(appEnv.DEFAULT_AGENT_CONFIG) || {};
19+
20+
await sessionModel.createInbox(defaultAgentConfig);
21+
}
22+
}

src/server/services/user/index.test.ts

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
44
import { UserItem } from '@/database/schemas';
55
import { UserModel } from '@/database/server/models/user';
66
import { pino } from '@/libs/logger';
7+
import { AgentService } from '@/server/services/agent';
78

89
import { UserService } from './index';
910

@@ -29,6 +30,12 @@ vi.mock('@/libs/logger', () => ({
2930
},
3031
}));
3132

33+
vi.mock('@/server/services/agent', () => ({
34+
AgentService: vi.fn().mockImplementation(() => ({
35+
createInbox: vi.fn().mockResolvedValue(undefined),
36+
})),
37+
}));
38+
3239
let service: UserService;
3340
const mockUserId = 'test-user-id';
3441

@@ -57,7 +64,7 @@ describe('UserService', () => {
5764
// Mock user not found
5865
vi.mocked(UserModel.findById).mockResolvedValue(null as any);
5966

60-
await service.createUser(mockUserId, mockUserJSON);
67+
const result = await service.createUser(mockUserId, mockUserJSON);
6168

6269
expect(UserModel.findById).toHaveBeenCalledWith(expect.anything(), mockUserId);
6370
expect(UserModel.createUser).toHaveBeenCalledWith(
@@ -73,6 +80,12 @@ describe('UserService', () => {
7380
clerkCreatedAt: new Date('2023-01-01T00:00:00Z'),
7481
}),
7582
);
83+
expect(AgentService).toHaveBeenCalledWith(expect.anything(), mockUserId);
84+
expect(vi.mocked(AgentService).mock.results[0].value.createInbox).toHaveBeenCalled();
85+
expect(result).toEqual({
86+
message: 'user created',
87+
success: true,
88+
});
7689
});
7790

7891
it('should not create user if already exists', async () => {
@@ -83,6 +96,7 @@ describe('UserService', () => {
8396

8497
expect(UserModel.findById).toHaveBeenCalledWith(expect.anything(), mockUserId);
8598
expect(UserModel.createUser).not.toHaveBeenCalled();
99+
expect(AgentService).not.toHaveBeenCalled();
86100
expect(result).toEqual({
87101
message: 'user not created due to user already existing in the database',
88102
success: false,
@@ -106,6 +120,8 @@ describe('UserService', () => {
106120
phone: '+1234567890', // Should use first phone number
107121
}),
108122
);
123+
expect(AgentService).toHaveBeenCalledWith(expect.anything(), mockUserId);
124+
expect(vi.mocked(AgentService).mock.results[0].value.createInbox).toHaveBeenCalled();
109125
});
110126
});
111127

src/server/services/user/index.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { serverDB } from '@/database/server';
44
import { UserModel } from '@/database/server/models/user';
55
import { pino } from '@/libs/logger';
66
import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
7+
import { AgentService } from '@/server/services/agent';
78

89
export class UserService {
910
createUser = async (id: string, params: UserJSON) => {
@@ -41,6 +42,10 @@ export class UserService {
4142
username: params.username,
4243
});
4344

45+
// 3. Create an inbox session for the user
46+
const agentService = new AgentService(serverDB, id);
47+
await agentService.createInbox();
48+
4449
/* ↓ cloud slot ↓ */
4550

4651
/* ↑ cloud slot ↑ */

0 commit comments

Comments
 (0)