Skip to content

refactor(ai): optimize AI session manager and provider factory #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ZeR020 merged 3 commits into from
May 7, 2026
Merged

refactor(ai): optimize AI session manager and provider factory #14

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ca0d0ff
refactor(ai): prepared statement caching in session manager, factory …
ZeR020 May 7, 2026
18ee9b3
fix(ai): make message sequence writes atomic
ZeR020 May 7, 2026
1539dee
fix(ai): make message sequence index migration safe
ZeR020 May 7, 2026
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
17 changes: 13 additions & 4 deletions src/services/ai/ai-provider-factory.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,14 @@ import type { AIProviderType } from "./session/session-types.js";

export class AIProviderFactory {
private static cleanupTimer: NodeJS.Timeout | null = null;
private static sessionManager: ReturnType<typeof getAISessionManager> | null = null;

private static getSessionManager() {
if (!this.sessionManager) {
this.sessionManager = getAISessionManager();
}
return this.sessionManager;
}

static startCleanupSchedule(intervalMs: number = 1000 * 60 * 60) {
if (this.cleanupTimer) clearInterval(this.cleanupTimer);
Expand All @@ -22,18 +30,19 @@ export class AIProviderFactory {
}

static createProvider(providerType: AIProviderType, config: ProviderConfig): BaseAIProvider {
const sessionManager = this.getSessionManager();
switch (providerType) {
case "openai-chat":
return new OpenAIChatCompletionProvider(config, getAISessionManager());
return new OpenAIChatCompletionProvider(config, sessionManager);

case "openai-responses":
return new OpenAIResponsesProvider(config, getAISessionManager());
return new OpenAIResponsesProvider(config, sessionManager);

case "anthropic":
return new AnthropicMessagesProvider(config, getAISessionManager());
return new AnthropicMessagesProvider(config, sessionManager);

case "google-gemini":
return new GoogleGeminiProvider(config, getAISessionManager());
return new GoogleGeminiProvider(config, sessionManager);

default:
throw new Error(`Unknown provider type: ${providerType}`);
Expand Down
172 changes: 86 additions & 86 deletions src/services/ai/session/ai-session-manager.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,18 @@ export class AISessionManager {
private db: DatabaseType;
private readonly dbPath: string;
private readonly sessionRetentionMs: number;
private getSessionStmt: any;
private getMessagesStmt: any;
private getLastSequenceStmt: any;
private cleanupExpiredStmt: any;
private deleteSessionStmt: any;
private getMessagesByRoleStmt: any;
private addMessageStmt: any;
private getNextSeqStmt: any;
private clearMessagesStmt: any;
private updateConversationIdStmt: any;
private updateMetadataStmt: any;
private updateBothStmt: any;

constructor() {
this.dbPath = join(CONFIG.storagePath, AI_SESSIONS_DB_NAME);
Expand All @@ -30,6 +42,40 @@ export class AISessionManager {
this.db = connectionManager.getConnection(this.dbPath);
this.sessionRetentionMs = CONFIG.aiSessionRetentionDays * 24 * 60 * 60 * 1000;
this.initDatabase();
this.getSessionStmt = this.db.prepare(`
SELECT * FROM ai_sessions WHERE session_id = ? AND provider = ? AND expires_at > ?
`);
this.getMessagesStmt = this.db.prepare(
"SELECT * FROM ai_messages WHERE ai_session_id = ? ORDER BY sequence ASC"
);
this.getLastSequenceStmt = this.db.prepare(
"SELECT MAX(sequence) as max_seq FROM ai_messages WHERE ai_session_id = ?"
);
this.cleanupExpiredStmt = this.db.prepare("DELETE FROM ai_sessions WHERE expires_at < ?");
this.deleteSessionStmt = this.db.prepare(
"DELETE FROM ai_sessions WHERE session_id = ? AND provider = ?"
);
this.getMessagesByRoleStmt = this.db.prepare(
"SELECT * FROM ai_messages WHERE ai_session_id = ? AND role = ? ORDER BY sequence ASC"
);
this.addMessageStmt = this.db.prepare(`
INSERT INTO ai_messages (
ai_session_id, sequence, role, content, tool_calls, tool_call_id, content_blocks, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
`);
this.getNextSeqStmt = this.db.prepare(
"SELECT COALESCE(MAX(sequence), -1) + 1 as next_seq FROM ai_messages WHERE ai_session_id = ?"
);
this.clearMessagesStmt = this.db.prepare("DELETE FROM ai_messages WHERE ai_session_id = ?");
this.updateConversationIdStmt = this.db.prepare(
"UPDATE ai_sessions SET conversation_id = ?, updated_at = ? WHERE session_id = ? AND provider = ?"
);
this.updateMetadataStmt = this.db.prepare(
"UPDATE ai_sessions SET metadata = ?, updated_at = ? WHERE session_id = ? AND provider = ?"
);
this.updateBothStmt = this.db.prepare(
"UPDATE ai_sessions SET conversation_id = ?, metadata = ?, updated_at = ? WHERE session_id = ? AND provider = ?"
);
}

private initDatabase(): void {
Expand Down Expand Up @@ -74,14 +120,8 @@ export class AISessionManager {
}

getSession(sessionId: string, provider: AIProviderType): AISession | null {
const stmt = this.db.prepare(`
SELECT * FROM ai_sessions
WHERE session_id = ? AND provider = ? AND expires_at > ?
`);
const row = stmt.get(sessionId, provider, Date.now()) as any;

const row = this.getSessionStmt.get(sessionId, provider, Date.now()) as any;
if (!row) return null;

return this.rowToSession(row);
}

Expand All @@ -93,7 +133,7 @@ export class AISessionManager {
this.db.run(
`
INSERT INTO ai_sessions (
id, provider, session_id, conversation_id,
id, provider, session_id, conversation_id,
metadata, created_at, updated_at, expires_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
`,
Expand All @@ -113,113 +153,73 @@ export class AISessionManager {
}

updateSession(sessionId: string, provider: AIProviderType, updates: SessionUpdateParams): void {
const fields: string[] = [];
const values: any[] = [];

if (updates.conversationId !== undefined) {
fields.push("conversation_id = ?");
values.push(updates.conversationId);
}

if (updates.metadata !== undefined) {
fields.push("metadata = ?");
values.push(JSON.stringify(updates.metadata));
const now = Date.now();
if (updates.conversationId !== undefined && updates.metadata !== undefined) {
this.updateBothStmt.run(
updates.conversationId,
JSON.stringify(updates.metadata),
now,
sessionId,
provider
);
} else if (updates.conversationId !== undefined) {
this.updateConversationIdStmt.run(updates.conversationId, now, sessionId, provider);
} else if (updates.metadata !== undefined) {
this.updateMetadataStmt.run(JSON.stringify(updates.metadata), now, sessionId, provider);
}

fields.push("updated_at = ?");
values.push(Date.now());

values.push(sessionId);
values.push(provider);

this.db.run(
`
UPDATE ai_sessions
SET ${fields.join(", ")}
WHERE session_id = ? AND provider = ?
`,
values
);
}

cleanupExpiredSessions(): number {
const result = this.db.run(`DELETE FROM ai_sessions WHERE expires_at < ?`, [Date.now()]);
const result = this.cleanupExpiredStmt.run(Date.now());
return result.changes;
}

deleteSession(sessionId: string, provider: AIProviderType): void {
this.db.run(`DELETE FROM ai_sessions WHERE session_id = ? AND provider = ?`, [
sessionId,
provider,
]);
this.deleteSessionStmt.run(sessionId, provider);
}

addMessage(message: Omit<AIMessage, "id" | "createdAt">): void {
this.db.run(
`INSERT INTO ai_messages (
ai_session_id, sequence, role, content,
tool_calls, tool_call_id, content_blocks, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
[
message.aiSessionId,
message.sequence,
message.role,
message.content,
message.toolCalls ? JSON.stringify(message.toolCalls) : null,
message.toolCallId || null,
message.contentBlocks ? JSON.stringify(message.contentBlocks) : null,
Date.now(),
]
this.addMessageStmt.run(
message.aiSessionId,
message.sequence,
message.role,
message.content,
message.toolCalls ? JSON.stringify(message.toolCalls) : null,
message.toolCallId || null,
message.contentBlocks ? JSON.stringify(message.contentBlocks) : null,
Date.now()
);
}

getMessages(aiSessionId: string): AIMessage[] {
const stmt = this.db.prepare(
"SELECT * FROM ai_messages WHERE ai_session_id = ? ORDER BY sequence ASC"
);
const rows = stmt.all(aiSessionId) as any[];

return rows.map(this.rowToMessage);
const rows = this.getMessagesStmt.all(aiSessionId) as any[];
return rows.map((row) => this.rowToMessage(row));
}

getLastSequence(aiSessionId: string): number {
const stmt = this.db.prepare(
"SELECT MAX(sequence) as max_seq FROM ai_messages WHERE ai_session_id = ?"
);
const row = stmt.get(aiSessionId) as any;

const row = this.getLastSequenceStmt.get(aiSessionId) as any;
return row?.max_seq ?? -1;
}

addMessageAtomic(message: Omit<AIMessage, "id" | "sequence" | "createdAt">): number {
const nextSeq = this.db
.prepare(
"SELECT COALESCE(MAX(sequence), -1) + 1 as next_seq FROM ai_messages WHERE ai_session_id = ?"
)
.get(message.aiSessionId) as { next_seq: number };
const nextSeq = this.getNextSeqStmt.get(message.aiSessionId) as { next_seq: number };

const seq = nextSeq.next_seq;
this.db.run(
`INSERT INTO ai_messages (
ai_session_id, sequence, role, content,
tool_calls, tool_call_id, content_blocks, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
[
message.aiSessionId,
seq,
message.role,
message.content,
message.toolCalls ? JSON.stringify(message.toolCalls) : null,
message.toolCallId || null,
message.contentBlocks ? JSON.stringify(message.contentBlocks) : null,
Date.now(),
]
this.addMessageStmt.run(
message.aiSessionId,
seq,
message.role,
message.content,
message.toolCalls ? JSON.stringify(message.toolCalls) : null,
message.toolCallId || null,
message.contentBlocks ? JSON.stringify(message.contentBlocks) : null,
Date.now()
);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
return seq;
}

clearMessages(aiSessionId: string): void {
this.db.run("DELETE FROM ai_messages WHERE ai_session_id = ?", [aiSessionId]);
this.clearMessagesStmt.run(aiSessionId);
}

private rowToSession(row: any): AISession {
Expand Down
Loading