Jayden/AFFiNE-Mirror
1
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-02-13 12:55:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(server): add cron job for session cleanup (#13181)

Browse Source 
fix AI-338
This commit is contained in:
DarkSky
2025-07-13 21:53:38 +08:00
committed by GitHub
parent 3ee82bd9ce
commit b6187718ea
11 changed files with 508 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.md
View File

@@ -372,3 +372,68 @@ Generated by [AVA](https://avajs.dev).
[assistant]: Quantum computing uses quantum mechanics principles.`,
promptName: 'Summary as title',
}
## should handle copilot cron jobs correctly
> daily job scheduling calls
[
{
args: [
'copilot.session.cleanupEmptySessions',
{},
{
jobId: 'daily-copilot-cleanup-empty-sessions',
},
],
},
{
args: [
'copilot.session.generateMissingTitles',
{},
{
jobId: 'daily-copilot-generate-missing-titles',
},
],
},
]
> cleanup empty sessions calls
[
{
args: [
'Date',
],
},
]
> title generation calls
{
jobCalls: [
{
args: [
'copilot.session.generateTitle',
{
sessionId: 'session1',
},
],
},
{
args: [
'copilot.session.generateTitle',
{
sessionId: 'session2',
},
],
},
],
modelCalls: [
{
args: [
100,
],
},
],
}

BIN
packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.snap
View File

Binary file not shown.

18
packages/backend/server/src/__tests__/copilot-provider.spec.ts
View File

@@ -351,10 +351,10 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
params: {
files: [
{
blobId: 'euclidean_distance',
fileName: 'euclidean_distance.rs',
fileType: 'text/rust',
fileContent: TestAssets.Code,
blobId: 'todo_md',
fileName: 'todo.md',
fileType: 'text/markdown',
fileContent: TestAssets.TODO,
},
],
},
@@ -476,6 +476,7 @@ The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Ca
},
},
],
config: { model: 'gemini-2.5-pro' },
verifier: (t: ExecutionContext<Tester>, result: string) => {
t.notThrows(() => {
TranscriptionResponseSchema.parse(JSON.parse(result));
@@ -697,11 +698,12 @@ for (const {
t.truthy(provider, 'should have provider');
await retry(`action: ${promptName}`, t, async t => {
const finalConfig = Object.assign({}, prompt.config, config);
const modelId = finalConfig.model || prompt.model;
switch (type) {
case 'text': {
const result = await provider.text(
{ modelId: prompt.model },
{ modelId },
[
...prompt.finish(
messages.reduce(
@@ -720,7 +722,7 @@ for (const {
}
case 'structured': {
const result = await provider.structure(
{ modelId: prompt.model },
{ modelId },
[
...prompt.finish(
messages.reduce(
@@ -739,7 +741,7 @@ for (const {
case 'object': {
const streamObjects: StreamObject[] = [];
for await (const chunk of provider.streamObject(
{ modelId: prompt.model },
{ modelId },
[
...prompt.finish(
messages.reduce(
@@ -771,7 +773,7 @@ for (const {
});
}
const stream = provider.streamImages(
{ modelId: prompt.model },
{ modelId },
[
...prompt.finish(
finalMessage.reduce(

72
packages/backend/server/src/__tests__/copilot.spec.ts
View File

@@ -18,6 +18,7 @@ import {
} from '../models';
import { CopilotModule } from '../plugins/copilot';
import { CopilotContextService } from '../plugins/copilot/context';
import { CopilotCronJobs } from '../plugins/copilot/cron';
import {
CopilotEmbeddingJob,
MockEmbeddingClient,
@@ -77,6 +78,7 @@ type Context = {
jobs: CopilotEmbeddingJob;
storage: CopilotStorage;
workflow: CopilotWorkflowService;
cronJobs: CopilotCronJobs;
executors: {
image: CopilotChatImageExecutor;
text: CopilotChatTextExecutor;
@@ -137,6 +139,7 @@ test.before(async t => {
const jobs = module.get(CopilotEmbeddingJob);
const transcript = module.get(CopilotTranscriptionService);
const workspaceEmbedding = module.get(CopilotWorkspaceService);
const cronJobs = module.get(CopilotCronJobs);
t.context.module = module;
t.context.auth = auth;
@@ -153,6 +156,7 @@ test.before(async t => {
t.context.jobs = jobs;
t.context.transcript = transcript;
t.context.workspaceEmbedding = workspaceEmbedding;
t.context.cronJobs = cronJobs;
t.context.executors = {
image: module.get(CopilotChatImageExecutor),
@@ -1931,3 +1935,71 @@ test('should handle generateSessionTitle correctly under various conditions', as
);
}
});
test('should handle copilot cron jobs correctly', async t => {
const { cronJobs, copilotSession } = t.context;
// mock calls
const mockCleanupResult = { removed: 2, cleaned: 3 };
const mockSessions = [
{ id: 'session1', _count: { messages: 1 } },
{ id: 'session2', _count: { messages: 2 } },
];
const cleanupStub = Sinon.stub(
copilotSession,
'cleanupEmptySessions'
).resolves(mockCleanupResult);
const toBeGenerateStub = Sinon.stub(
copilotSession,
'toBeGenerateTitle'
).resolves(mockSessions);
const jobAddStub = Sinon.stub(cronJobs['jobs'], 'add').resolves();
// daily cleanup job scheduling
{
await cronJobs.dailyCleanupJob();
t.snapshot(
jobAddStub.getCalls().map(call => ({
args: call.args,
})),
'daily job scheduling calls'
);
jobAddStub.reset();
cleanupStub.reset();
toBeGenerateStub.reset();
}
// cleanup empty sessions
{
// mock
cleanupStub.resolves(mockCleanupResult);
toBeGenerateStub.resolves(mockSessions);
await cronJobs.cleanupEmptySessions();
t.snapshot(
cleanupStub.getCalls().map(call => ({
args: call.args.map(arg => (arg instanceof Date ? 'Date' : arg)), // Replace Date with string for stable snapshot
})),
'cleanup empty sessions calls'
);
}
// generate missing titles
await cronJobs.generateMissingTitles();
t.snapshot(
{
modelCalls: toBeGenerateStub.getCalls().map(call => ({
args: call.args,
})),
jobCalls: jobAddStub.getCalls().map(call => ({
args: call.args,
})),
},
'title generation calls'
);
cleanupStub.restore();
toBeGenerateStub.restore();
jobAddStub.restore();
});

62
packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.md
View File

@@ -565,3 +565,65 @@ Generated by [AVA](https://avajs.dev).
workspaceSessionExists: true,
},
}
## should cleanup empty sessions correctly
> cleanup empty sessions results
{
cleanupResult: {
cleaned: 0,
removed: 0,
},
remainingSessions: [
{
deleted: false,
pinned: false,
type: 'zeroCost',
},
{
deleted: false,
pinned: false,
type: 'zeroCost',
},
{
deleted: false,
pinned: false,
type: 'noMessages',
},
{
deleted: false,
pinned: false,
type: 'noMessages',
},
{
deleted: false,
pinned: false,
type: 'recent',
},
{
deleted: false,
pinned: false,
type: 'withMessages',
},
],
}
## should get sessions for title generation correctly
> sessions for title generation results
{
onlyValidSessionsReturned: true,
sessions: [
{
assistantMessageCount: 1,
isValid: true,
},
{
assistantMessageCount: 2,
isValid: true,
},
],
total: 2,
}

BIN
packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.snap
View File

Binary file not shown.

175
packages/backend/server/src/__tests__/models/copilot-session.spec.ts
View File

@@ -917,3 +917,178 @@ test('should handle fork and session attachment operations', async t => {
'attach and detach operation results'
);
});
test('should cleanup empty sessions correctly', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const oneDayAgo = new Date(Date.now() - 24 * 60 * 60 * 1000);
const twoHoursAgo = new Date(Date.now() - 2 * 60 * 60 * 1000);
// should be deleted
const neverUsedSessionIds: string[] = [randomUUID(), randomUUID()];
await Promise.all(
neverUsedSessionIds.map(async id => {
await createTestSession(t, { sessionId: id });
await db.aiSession.update({
where: { id },
data: { messageCost: 0, updatedAt: oneDayAgo },
});
})
);
// should be marked as deleted
const emptySessionIds: string[] = [randomUUID(), randomUUID()];
await Promise.all(
emptySessionIds.map(async id => {
await createTestSession(t, { sessionId: id });
await db.aiSession.update({
where: { id },
data: { messageCost: 100, updatedAt: oneDayAgo },
});
})
);
// should not be affected
const recentSessionId = randomUUID();
await createTestSession(t, { sessionId: recentSessionId });
await db.aiSession.update({
where: { id: recentSessionId },
data: { messageCost: 0, updatedAt: twoHoursAgo },
});
// Create session with messages (should not be affected)
const sessionWithMsgId = randomUUID();
await createSessionWithMessages(
t,
{ sessionId: sessionWithMsgId },
'test message'
);
const result = await copilotSession.cleanupEmptySessions(oneDayAgo);
const remainingSessions = await db.aiSession.findMany({
where: {
id: {
in: [
...neverUsedSessionIds,
...emptySessionIds,
recentSessionId,
sessionWithMsgId,
],
},
},
select: { id: true, deletedAt: true, pinned: true },
});
t.snapshot(
{
cleanupResult: result,
remainingSessions: remainingSessions.map(s => ({
deleted: !!s.deletedAt,
pinned: s.pinned,
type: neverUsedSessionIds.includes(s.id)
? 'zeroCost'
: emptySessionIds.includes(s.id)
? 'noMessages'
: s.id === recentSessionId
? 'recent'
: 'withMessages',
})),
},
'cleanup empty sessions results'
);
});
test('should get sessions for title generation correctly', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
// create valid sessions with messages
const sessionIds: string[] = [randomUUID(), randomUUID()];
await Promise.all(
sessionIds.map(async (id, index) => {
await createTestSession(t, { sessionId: id });
await db.aiSession.update({
where: { id },
data: {
updatedAt: new Date(Date.now() - index * 1000),
messages: {
create: Array.from({ length: index + 1 }, (_, i) => ({
role: 'assistant',
content: `assistant message ${i}`,
})),
},
},
});
})
);
// create excluded sessions
const excludedSessions = [
{
reason: 'hasTitle',
setupFn: async (id: string) => {
await createTestSession(t, { sessionId: id });
await db.aiSession.update({
where: { id },
data: { title: 'Existing Title' },
});
},
},
{
reason: 'isDeleted',
setupFn: async (id: string) => {
await createTestSession(t, { sessionId: id });
await db.aiSession.update({
where: { id },
data: { deletedAt: new Date() },
});
},
},
{
reason: 'noMessages',
setupFn: async (id: string) => {
await createTestSession(t, { sessionId: id });
},
},
{
reason: 'isAction',
setupFn: async (id: string) => {
await createTestSession(t, {
sessionId: id,
promptName: TEST_PROMPTS.ACTION,
});
},
},
{
reason: 'noAssistantMessages',
setupFn: async (id: string) => {
await createTestSession(t, { sessionId: id });
await db.aiSessionMessage.create({
data: { sessionId: id, role: 'user', content: 'User message only' },
});
},
},
];
await Promise.all(
excludedSessions.map(async session => {
await session.setupFn(randomUUID());
})
);
const result = await copilotSession.toBeGenerateTitle(10);
t.snapshot(
{
total: result.length,
sessions: result.map(s => ({
assistantMessageCount: s._count.messages,
isValid: sessionIds.includes(s.id),
})),
onlyValidSessionsReturned: result.every(s => sessionIds.includes(s.id)),
},
'sessions for title generation results'
);
});

53
packages/backend/server/src/models/copilot-session.ts
View File

@@ -582,4 +582,57 @@ export class CopilotSessionModel extends BaseModel {
.map(({ messageCost, prompt: { action } }) => (action ? 1 : messageCost))
.reduce((prev, cost) => prev + cost, 0);
}
@Transactional()
async cleanupEmptySessions(earlyThen: Date) {
// delete never used sessions
const { count: removed } = await this.db.aiSession.deleteMany({
where: {
messageCost: 0,
deletedAt: null,
// filter session updated more than 24 hours ago
updatedAt: { lt: earlyThen },
},
});
// mark empty sessions as deleted
const { count: cleaned } = await this.db.aiSession.updateMany({
where: {
deletedAt: null,
messages: { none: {} },
// filter session updated more than 24 hours ago
updatedAt: { lt: earlyThen },
},
data: {
deletedAt: new Date(),
pinned: false,
},
});
return { removed, cleaned };
}
@Transactional()
async toBeGenerateTitle(take: number) {
const sessions = await this.db.aiSession
.findMany({
where: {
title: null,
deletedAt: null,
messages: { some: {} },
// only generate titles for non-actions sessions
prompt: { action: null },
},
select: {
id: true,
// count assistant messages
_count: { select: { messages: { where: { role: 'assistant' } } } },
},
take,
orderBy: { updatedAt: 'desc' },
})
.then(s => s.filter(s => s._count.messages > 0));
return sessions;
}
}

67
packages/backend/server/src/plugins/copilot/cron.ts Normal file
View File

@@ -0,0 +1,67 @@
import { Injectable, Logger } from '@nestjs/common';
import { Cron, CronExpression } from '@nestjs/schedule';
import { JobQueue, OneDay, OnJob } from '../../base';
import { Models } from '../../models';
declare global {
interface Jobs {
'copilot.session.cleanupEmptySessions': {};
'copilot.session.generateMissingTitles': {};
}
}
const GENERATE_TITLES_BATCH_SIZE = 100;
@Injectable()
export class CopilotCronJobs {
private readonly logger = new Logger(CopilotCronJobs.name);
constructor(
private readonly models: Models,
private readonly jobs: JobQueue
) {}
@Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT)
async dailyCleanupJob() {
await this.jobs.add(
'copilot.session.cleanupEmptySessions',
{},
{ jobId: 'daily-copilot-cleanup-empty-sessions' }
);
await this.jobs.add(
'copilot.session.generateMissingTitles',
{},
{ jobId: 'daily-copilot-generate-missing-titles' }
);
}
@OnJob('copilot.session.cleanupEmptySessions')
async cleanupEmptySessions() {
const { removed, cleaned } =
await this.models.copilotSession.cleanupEmptySessions(
new Date(Date.now() - OneDay)
);
this.logger.log(
`Cleanup completed: ${removed} sessions deleted, ${cleaned} sessions marked as deleted`
);
}
@OnJob('copilot.session.generateMissingTitles')
async generateMissingTitles() {
const sessions = await this.models.copilotSession.toBeGenerateTitle(
GENERATE_TITLES_BATCH_SIZE
);
for (const session of sessions) {
await this.jobs.add('copilot.session.generateTitle', {
sessionId: session.id,
});
}
this.logger.log(
`Scheduled title generation for ${sessions.length} sessions`
);
}
}

3
packages/backend/server/src/plugins/copilot/index.ts
View File

@@ -15,6 +15,7 @@ import {
CopilotContextService,
} from './context';
import { CopilotController } from './controller';
import { CopilotCronJobs } from './cron';
import { CopilotEmbeddingJob } from './embedding';
import { ChatMessageCache } from './message';
import { PromptService } from './prompt';
@@ -64,6 +65,8 @@ import {
CopilotContextResolver,
CopilotContextService,
CopilotEmbeddingJob,
// cron jobs
CopilotCronJobs,
// transcription
CopilotTranscriptionService,
CopilotTranscriptionResolver,

1
packages/backend/server/src/plugins/copilot/prompt/prompts.ts
View File

@@ -304,6 +304,7 @@ const textActions: Prompt[] = [
name: 'Transcript audio',
action: 'Transcript audio',
model: 'gemini-2.5-flash',
optionalModels: ['gemini-2.5-flash', 'gemini-2.5-pro'],
messages: [
{
role: 'system',