-
Notifications
You must be signed in to change notification settings - Fork 2
fix: harden memory services and web UI #9
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.
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
Changes from 4 commits
7d73f37
71b053f
202232e
69fa389
093e029
4c2aa9a
e9d5dfa
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -39,4 +39,3 @@ function main() { | |
| } | ||
|
|
||
| main(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -61,4 +61,3 @@ function migrateDirectory(dir) { | |
| } | ||
|
|
||
| migrateDirectory("tests"); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -405,4 +405,3 @@ migrate(storagePath) | |
| console.error("Migration failed:", error); | ||
| process.exit(1); | ||
| }); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -48,4 +48,3 @@ export class AIProviderFactory { | |
| return getAISessionManager().cleanupExpiredSessions(); | ||
| } | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -297,4 +297,3 @@ export async function generateStructuredOutput<T>(options: { | |
| }); | ||
| return result.output as T; | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -237,4 +237,3 @@ export class AnthropicMessagesProvider extends BaseAIProvider { | |
| return null; | ||
| } | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -61,4 +61,3 @@ export abstract class BaseAIProvider { | |
|
|
||
| abstract supportsSession(): boolean; | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -434,4 +434,3 @@ export class OpenAIChatCompletionProvider extends BaseAIProvider { | |
| }; | ||
| } | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -233,4 +233,3 @@ export class OpenAIResponsesProvider extends BaseAIProvider { | |
| return data; | ||
| } | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -62,4 +62,3 @@ export class ToolSchemaConverter { | |
| }; | ||
| } | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -118,4 +118,3 @@ export class UserProfileValidator { | |
| return errors; | ||
| } | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -142,12 +142,15 @@ export async function handleListMemories( | |
| includePrompts: boolean = true | ||
| ): Promise<ApiResponse<PaginatedResponse<Memory | any>>> { | ||
| try { | ||
| const safePage = Number.isFinite(page) && page > 0 ? Math.floor(page) : 1; | ||
| const safePageSize = | ||
| Number.isFinite(pageSize) && pageSize > 0 && pageSize <= 100 ? Math.floor(pageSize) : 20; | ||
| await embeddingService.warmup(); | ||
| let allMemories: any[] = []; | ||
| if (tag) { | ||
| const { scope: tagScope, hash } = extractScopeFromTag(tag); | ||
| const shards = shardManager.getAllShards(tagScope, hash); | ||
| const limit = page * pageSize; // Fetch enough to cover the requested page across shards | ||
| const limit = safePage * safePageSize; // Fetch enough to cover the requested page across shards | ||
| for (const shard of shards) { | ||
| const db = connectionManager.getConnection(shard.dbPath); | ||
| const memories = vectorSearch.listMemories(db, tag, limit); | ||
|
|
@@ -233,9 +236,9 @@ export async function handleListMemories( | |
| timeline = sortedTimeline; | ||
|
|
||
| const total = timeline.length; | ||
| const totalPages = Math.ceil(total / pageSize); | ||
| const offset = (page - 1) * pageSize; | ||
| const paginatedResults = timeline.slice(offset, offset + pageSize); | ||
| const totalPages = Math.ceil(total / safePageSize); | ||
| const offset = (safePage - 1) * safePageSize; | ||
| const paginatedResults = timeline.slice(offset, offset + safePageSize); | ||
|
|
||
| const items = paginatedResults.map((item: any) => { | ||
| if (item.type === "memory") { | ||
|
|
@@ -491,6 +494,9 @@ export async function handleSearch( | |
| ): Promise<ApiResponse<PaginatedResponse<SearchResultItem>>> { | ||
| try { | ||
| if (!query) return { success: false, error: "query is required" }; | ||
| const safePage = Number.isFinite(page) && page > 0 ? Math.floor(page) : 1; | ||
| const safePageSize = | ||
| Number.isFinite(pageSize) && pageSize > 0 && pageSize <= 100 ? Math.floor(pageSize) : 20; | ||
| await embeddingService.warmup(); | ||
| const queryVector = await embeddingService.embedWithTimeout(query); | ||
| let memoryResults: any[] = []; | ||
|
|
@@ -500,39 +506,32 @@ export async function handleSearch( | |
| const shards = shardManager.getAllShards(scope, hash); | ||
| for (const shard of shards) { | ||
| try { | ||
| const results = await vectorSearch.searchInShard(shard, queryVector, tag, pageSize * 2); | ||
| const results = await vectorSearch.searchInShard( | ||
| shard, | ||
| queryVector, | ||
| tag, | ||
| safePageSize * 2 | ||
| ); | ||
| memoryResults.push(...results); | ||
| } catch (error) { | ||
| log("Shard search error", { shardId: shard.id, error: String(error) }); | ||
| } | ||
| } | ||
| const projectPath = getProjectPathFromTag(tag); | ||
| promptResults = userPromptManager.searchPrompts(query, projectPath, pageSize * 2); | ||
| promptResults = userPromptManager.searchPrompts(query, projectPath, safePageSize * 2); | ||
| } else { | ||
| const userShards = shardManager.getAllShards("user", ""); | ||
| const projectShards = shardManager.getAllShards("project", ""); | ||
| const uniqueTags = new Set<string>(); | ||
| for (const shard of projectShards) { | ||
| const db = connectionManager.getConnection(shard.dbPath); | ||
| const tags = vectorSearch.getDistinctTags(db); | ||
| for (const t of tags) { | ||
| if (t.container_tag) uniqueTags.add(t.container_tag); | ||
| } | ||
| } | ||
| for (const containerTag of uniqueTags) { | ||
| const { scope, hash } = extractScopeFromTag(containerTag); | ||
| const shards = shardManager.getAllShards(scope, hash); | ||
| for (const shard of shards) { | ||
| try { | ||
| const results = await vectorSearch.searchInShard( | ||
| shard, | ||
| queryVector, | ||
| containerTag, | ||
| pageSize | ||
| ); | ||
| memoryResults.push(...results); | ||
| } catch (error) { | ||
| log("Shard search error", { shardId: shard.id, error: String(error) }); | ||
| } | ||
| const searchedPaths = new Set<string>(); | ||
| for (const shard of [...userShards, ...projectShards]) { | ||
|
Comment on lines
+526
to
+529
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When the web UI searches with no selected tag, it calls Useful? React with 👍 / 👎. |
||
| if (searchedPaths.has(shard.dbPath)) continue; | ||
| searchedPaths.add(shard.dbPath); | ||
| try { | ||
| const perShardLimit = Math.min(safePage * safePageSize, 500); | ||
| const results = await vectorSearch.searchInShard(shard, queryVector, "", perShardLimit); | ||
| memoryResults.push(...results); | ||
|
coderabbitai[bot] marked this conversation as resolved.
|
||
| } catch (error) { | ||
| log("Shard search error", { shardId: shard.id, error: String(error) }); | ||
| } | ||
| } | ||
| promptResults = userPromptManager.searchPrompts(query, undefined, pageSize * 2); | ||
|
|
@@ -575,9 +574,12 @@ export async function handleSearch( | |
| ); | ||
|
|
||
| const total = combinedResults.length; | ||
| const totalPages = Math.ceil(total / pageSize); | ||
| const offset = (page - 1) * pageSize; | ||
| const paginatedResults: SearchResultItem[] = combinedResults.slice(offset, offset + pageSize); | ||
| const totalPages = Math.ceil(total / safePageSize); | ||
| const offset = (safePage - 1) * safePageSize; | ||
| const paginatedResults: SearchResultItem[] = combinedResults.slice( | ||
| offset, | ||
| offset + safePageSize | ||
| ); | ||
|
|
||
| const missingPromptIds = new Set<string>(); | ||
| const missingMemoryIds = new Set<string>(); | ||
|
|
@@ -1179,4 +1181,3 @@ export async function handleConflictStats(): Promise< | |
| return { success: false, error: "Internal error" }; | ||
| } | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -131,4 +131,3 @@ export class CleanupService { | |
| } | ||
|
|
||
| export const cleanupService = new CleanupService(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -444,4 +444,3 @@ export class LocalMemoryClient { | |
| } | ||
|
|
||
| export const memoryClient = new LocalMemoryClient(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -40,4 +40,3 @@ export function formatContextForPrompt( | |
|
|
||
| return parts.join("\n"); | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -37,4 +37,3 @@ export function getLanguageName(code: string): string { | |
| } | ||
| return lang?.name || "English"; | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -619,4 +619,3 @@ export function recordAccess(currentAccessCount: number): { | |
| lastAccessed: Date.now(), | ||
| }; | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -350,4 +350,3 @@ export class MigrationService { | |
| } | ||
|
|
||
| export const migrationService = new MigrationService(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -164,4 +164,3 @@ export function calculateDiversityPenalty( | |
|
|
||
| return 0; | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -83,4 +83,3 @@ export class ConnectionManager { | |
| } | ||
|
|
||
| export const connectionManager = new ConnectionManager(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -405,4 +405,3 @@ export class ShardManager { | |
| } | ||
|
|
||
| export const shardManager = new ShardManager(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -126,4 +126,3 @@ export function getDatabase(): new (path: string) => Database { | |
| } | ||
| return DatabaseImpl; | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -275,4 +275,3 @@ export class TranscriptManager { | |
| } | ||
|
|
||
| export const transcriptManager = new TranscriptManager(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -81,4 +81,3 @@ export interface SearchResult { | |
| contextBoost?: number; | ||
| finalScore?: number; | ||
| } | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -675,4 +675,3 @@ export class VectorSearch { | |
| } | ||
|
|
||
| export const vectorSearch = new VectorSearch(); | ||
| // AUDIT_MARKER | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -178,4 +178,3 @@ export function getTags(directory: string): { | |
| project: getProjectTagInfo(directory), | ||
| }; | ||
| } | ||
| // AUDIT_MARKER | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Bound the page window before using it as a shard read limit.
safePageSizeis capped, butsafePageis not. On Line 153, a very largepagecan still turnlistMemoriesinto an effectively unbounded read and materialize far more rows than this endpoint can safely return. Cap or reject oversized page windows before calling storage.Suggested fix
try { const safePage = Number.isFinite(page) && page > 0 ? Math.floor(page) : 1; const safePageSize = Number.isFinite(pageSize) && pageSize > 0 && pageSize <= 100 ? Math.floor(pageSize) : 20; + const requestedWindow = safePage * safePageSize; + if (requestedWindow > 500) { + return { success: false, error: "Requested page window is too large" }; + } await embeddingService.warmup(); let allMemories: any[] = []; if (tag) { const { scope: tagScope, hash } = extractScopeFromTag(tag); const shards = shardManager.getAllShards(tagScope, hash); - const limit = safePage * safePageSize; // Fetch enough to cover the requested page across shards + const limit = requestedWindow;Also applies to: 153-153
🤖 Prompt for AI Agents