From 99b9bb847fe0f6d0542d52621cdf67fbbc47a683 Mon Sep 17 00:00:00 2001 From: Rohit Ghumare Date: Fri, 27 Feb 2026 11:57:34 +0530 Subject: [PATCH 1/3] fix: comprehensive security and reliability audit fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Critical: - Add auth to health, sessions, observations, MCP tools list endpoints - Remove ~/.claude from migration ALLOWED_DIRS (path traversal risk) - Add import payload size limits (10K sessions, 50K memories, 5K obs/session) - Fix index persistence: loaded indices now restore into active singletons High: - Cap expandIds to 20, BFS traversal to 500 nodes / 5 hops - Fetch relations once before BFS loop (was O(N²) KV reads) - Add observe input validation (sessionId, hookType, timestamp) - Sanitize provider error messages in compress and migrate - Tighten CSP: remove unsafe-inline Medium/Low: - Cap auto-forget contradiction comparison to 1000 memories - Add secret patterns for Anthropic, GitHub PAT, Gemini keys - Validate files/concepts arrays in remember - Wrap deserialize() in try/catch for corrupt data resilience - Fix parseInt NaN fallback in config - Remove dead code (ObservationQueue) --- src/config.ts | 12 +++---- src/functions/auto-forget.ts | 6 ++-- src/functions/compress.ts | 2 +- src/functions/export-import.ts | 32 +++++++++++++++++ src/functions/migrate.ts | 7 ++-- src/functions/observe.ts | 16 +++++++++ src/functions/privacy.ts | 28 +++++++++------ src/functions/relations.ts | 14 ++++---- src/functions/remember.ts | 12 ++++++- src/functions/smart-search.ts | 5 +-- src/health/recovery.ts | 22 ------------ src/index.ts | 21 ++++++----- src/mcp/server.ts | 12 ++++--- src/state/search-index.ts | 36 ++++++++++++------- src/state/vector-index.ts | 38 +++++++++++--------- src/triggers/api.ts | 65 ++++++++++++++++++++-------------- 16 files changed, 201 insertions(+), 127 deletions(-) delete mode 100644 src/health/recovery.ts diff --git a/src/config.ts b/src/config.ts index b9545a5..2858142 100644 --- a/src/config.ts +++ b/src/config.ts @@ -71,14 +71,12 @@ export function loadConfig(): AgentMemoryConfig { return { engineUrl: env["III_ENGINE_URL"] || "ws://localhost:49134", - restPort: parseInt(env["III_REST_PORT"] || "3111", 10), - streamsPort: parseInt(env["III_STREAMS_PORT"] || "3112", 10), + restPort: parseInt(env["III_REST_PORT"] || "3111", 10) || 3111, + streamsPort: parseInt(env["III_STREAMS_PORT"] || "3112", 10) || 3112, provider, - tokenBudget: parseInt(env["TOKEN_BUDGET"] || "2000", 10), - maxObservationsPerSession: parseInt( - env["MAX_OBS_PER_SESSION"] || "500", - 10, - ), + tokenBudget: parseInt(env["TOKEN_BUDGET"] || "2000", 10) || 2000, + maxObservationsPerSession: + parseInt(env["MAX_OBS_PER_SESSION"] || "500", 10) || 500, compressionModel: provider.model, dataDir: DATA_DIR, }; diff --git a/src/functions/auto-forget.ts b/src/functions/auto-forget.ts index 510c51a..96ba575 100644 --- a/src/functions/auto-forget.ts +++ b/src/functions/auto-forget.ts @@ -52,9 +52,9 @@ export function registerAutoForgetFunction(sdk: ISdk, kv: StateKV): void { } } - const latestMemories = memories.filter( - (m) => m.isLatest !== false && !deletedIds.has(m.id), - ); + const latestMemories = memories + .filter((m) => m.isLatest !== false && !deletedIds.has(m.id)) + .slice(0, 1000); for (let i = 0; i < latestMemories.length; i++) { for (let j = i + 1; j < latestMemories.length; j++) { const sim = jaccardSimilarity( diff --git a/src/functions/compress.ts b/src/functions/compress.ts index ead2775..8e8bf10 100644 --- a/src/functions/compress.ts +++ b/src/functions/compress.ts @@ -175,7 +175,7 @@ export function registerCompressFunction( obsId: data.observationId, error: msg, }); - return { success: false, error: msg }; + return { success: false, error: "compression_failed" }; } }, ); diff --git a/src/functions/export-import.ts b/src/functions/export-import.ts index 62b5a07..0fa0788 100644 --- a/src/functions/export-import.ts +++ b/src/functions/export-import.ts @@ -85,6 +85,38 @@ export function registerExportImportFunction(sdk: ISdk, kv: StateKV): void { }; } + const MAX_SESSIONS = 10_000; + const MAX_MEMORIES = 50_000; + const MAX_SUMMARIES = 10_000; + const MAX_OBS_PER_SESSION = 5_000; + + if (importData.sessions?.length > MAX_SESSIONS) { + return { + success: false, + error: `Too many sessions (max ${MAX_SESSIONS})`, + }; + } + if (importData.memories?.length > MAX_MEMORIES) { + return { + success: false, + error: `Too many memories (max ${MAX_MEMORIES})`, + }; + } + if (importData.summaries?.length > MAX_SUMMARIES) { + return { + success: false, + error: `Too many summaries (max ${MAX_SUMMARIES})`, + }; + } + for (const [, obs] of Object.entries(importData.observations || {})) { + if (obs.length > MAX_OBS_PER_SESSION) { + return { + success: false, + error: `Too many observations per session (max ${MAX_OBS_PER_SESSION})`, + }; + } + } + const stats = { sessions: 0, observations: 0, diff --git a/src/functions/migrate.ts b/src/functions/migrate.ts index f20c6bd..cb7f28a 100644 --- a/src/functions/migrate.ts +++ b/src/functions/migrate.ts @@ -10,10 +10,7 @@ import type { SessionSummary, } from "../types.js"; -const ALLOWED_DIRS = [ - resolve(homedir(), ".agentmemory"), - resolve(homedir(), ".claude"), -]; +const ALLOWED_DIRS = [resolve(homedir(), ".agentmemory")]; function isAllowedPath(dbPath: string): boolean { const resolved = resolve(dbPath); @@ -149,7 +146,7 @@ export function registerMigrateFunction(sdk: ISdk, kv: StateKV): void { } catch (err) { const msg = err instanceof Error ? err.message : String(err); ctx.logger.error("Migration failed", { error: msg }); - return { success: false, error: msg }; + return { success: false, error: "Migration failed" }; } }, ); diff --git a/src/functions/observe.ts b/src/functions/observe.ts index 389c0bf..1fa7e18 100644 --- a/src/functions/observe.ts +++ b/src/functions/observe.ts @@ -18,6 +18,22 @@ export function registerObserveFunction( }, async (payload: HookPayload) => { const ctx = getContext(); + + if ( + !payload?.sessionId || + typeof payload.sessionId !== "string" || + !payload.hookType || + typeof payload.hookType !== "string" || + !payload.timestamp || + typeof payload.timestamp !== "string" + ) { + return { + success: false, + error: + "Invalid payload: sessionId, hookType, and timestamp are required", + }; + } + const obsId = generateId("obs"); if (dedupMap) { diff --git a/src/functions/privacy.ts b/src/functions/privacy.ts index e93b506..3237356 100644 --- a/src/functions/privacy.ts +++ b/src/functions/privacy.ts @@ -1,30 +1,36 @@ -import type { ISdk } from 'iii-sdk' +import type { ISdk } from "iii-sdk"; -const PRIVATE_TAG_RE = /[\s\S]*?<\/private>/gi +const PRIVATE_TAG_RE = /[\s\S]*?<\/private>/gi; const SECRET_PATTERN_SOURCES = [ /(?:api[_-]?key|secret|token|password|credential|auth)[\s]*[=:]\s*["']?[A-Za-z0-9_\-/.+]{20,}["']?/gi, /(?:sk|pk|rk|ak)-[A-Za-z0-9]{20,}/g, + /sk-ant-[A-Za-z0-9\-_]{20,}/g, /ghp_[A-Za-z0-9]{36}/g, + /github_pat_[A-Za-z0-9_]{22,}/g, /xoxb-[A-Za-z0-9\-]+/g, /AKIA[0-9A-Z]{16}/g, + /AIza[A-Za-z0-9\-_]{35}/g, /eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}/g, -] +]; export function stripPrivateData(input: string): string { - let result = input.replace(PRIVATE_TAG_RE, '[REDACTED]') + let result = input.replace(PRIVATE_TAG_RE, "[REDACTED]"); for (const source of SECRET_PATTERN_SOURCES) { - const pattern = new RegExp(source.source, source.flags) - result = result.replace(pattern, '[REDACTED_SECRET]') + const pattern = new RegExp(source.source, source.flags); + result = result.replace(pattern, "[REDACTED_SECRET]"); } - return result + return result; } export function registerPrivacyFunction(sdk: ISdk): void { sdk.registerFunction( - { id: 'mem::privacy', description: 'Strip private tags and secrets from input' }, + { + id: "mem::privacy", + description: "Strip private tags and secrets from input", + }, async (data: { input: string }) => { - return { output: stripPrivateData(data.input) } - } - ) + return { output: stripPrivateData(data.input) }; + }, + ); } diff --git a/src/functions/relations.ts b/src/functions/relations.ts index 910c9ea..dcb6af7 100644 --- a/src/functions/relations.ts +++ b/src/functions/relations.ts @@ -112,7 +112,12 @@ export function registerRelationsFunction(sdk: ISdk, kv: StateKV): void { }, async (data: { memoryId: string; maxHops?: number }) => { const ctx = getContext(); - const maxHops = data.maxHops ?? 2; + const maxHops = Math.min(data.maxHops ?? 2, 5); + const MAX_VISITED = 500; + + const allRelations = await kv + .list(KV.relations) + .catch(() => []); const visited = new Set(); const result: Array<{ memory: Memory; hop: number }> = []; @@ -120,7 +125,7 @@ export function registerRelationsFunction(sdk: ISdk, kv: StateKV): void { { id: data.memoryId, hop: 0 }, ]; - while (queue.length > 0) { + while (queue.length > 0 && visited.size < MAX_VISITED) { const current = queue.shift()!; if (visited.has(current.id) || current.hop > maxHops) continue; visited.add(current.id); @@ -136,10 +141,7 @@ export function registerRelationsFunction(sdk: ISdk, kv: StateKV): void { const supersedes = memory.supersedes || []; const parentId = memory.parentId ? [memory.parentId] : []; - const kvRelations = await kv - .list(KV.relations) - .catch(() => []); - const kvLinked = kvRelations + const kvLinked = allRelations .filter((r) => r.sourceId === current.id || r.targetId === current.id) .map((r) => (r.sourceId === current.id ? r.targetId : r.sourceId)); diff --git a/src/functions/remember.ts b/src/functions/remember.ts index 739f633..1f42da7 100644 --- a/src/functions/remember.ts +++ b/src/functions/remember.ts @@ -14,9 +14,19 @@ export function registerRememberFunction(sdk: ISdk, kv: StateKV): void { files?: string[]; }) => { const ctx = getContext(); - if (!data.content || !data.content.trim()) { + if ( + !data.content || + typeof data.content !== "string" || + !data.content.trim() + ) { return { success: false, error: "content is required" }; } + if (data.files && !Array.isArray(data.files)) { + return { success: false, error: "files must be an array" }; + } + if (data.concepts && !Array.isArray(data.concepts)) { + return { success: false, error: "concepts must be an array" }; + } const validTypes = new Set([ "pattern", "preference", diff --git a/src/functions/smart-search.ts b/src/functions/smart-search.ts index 7c59c6e..79c59d4 100644 --- a/src/functions/smart-search.ts +++ b/src/functions/smart-search.ts @@ -23,13 +23,14 @@ export function registerSmartSearchFunction( const ctx = getContext(); if (data.expandIds && data.expandIds.length > 0) { + const ids = data.expandIds.slice(0, 20); const expanded: Array<{ obsId: string; sessionId: string; observation: CompressedObservation; }> = []; - for (const obsId of data.expandIds) { + for (const obsId of ids) { const obs = await findObservation(kv, obsId); if (obs) { expanded.push({ @@ -41,7 +42,7 @@ export function registerSmartSearchFunction( } ctx.logger.info("Smart search expanded", { - requested: data.expandIds.length, + requested: ids.length, found: expanded.length, }); return { mode: "expanded", results: expanded }; diff --git a/src/health/recovery.ts b/src/health/recovery.ts deleted file mode 100644 index 34be8bb..0000000 --- a/src/health/recovery.ts +++ /dev/null @@ -1,22 +0,0 @@ -import type { HookPayload } from "../types.js"; - -const MAX_QUEUE_SIZE = 500; - -export class ObservationQueue { - private queue: HookPayload[] = []; - - enqueue(payload: HookPayload): boolean { - if (this.queue.length >= MAX_QUEUE_SIZE) return false; - this.queue.push(payload); - return true; - } - - drain(): HookPayload[] { - const items = this.queue.splice(0); - return items; - } - - get size(): number { - return this.queue.length; - } -} diff --git a/src/index.ts b/src/index.ts index 90c120f..4bc6c5f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -140,18 +140,21 @@ async function main() { console.warn(`[agentmemory] Failed to load persisted index:`, err); return null; }); - if (loaded?.bm25) { - const restoredCount = loaded.bm25.size; - if (restoredCount > 0) { - console.log( - `[agentmemory] Loaded persisted BM25 index (${restoredCount} docs)`, - ); - } + if (loaded?.bm25 && loaded.bm25.size > 0) { + bm25Index.restoreFrom(loaded.bm25); + console.log( + `[agentmemory] Loaded persisted BM25 index (${bm25Index.size} docs)`, + ); + } + if (loaded?.vector && vectorIndex && loaded.vector.size > 0) { + vectorIndex.restoreFrom(loaded.vector); + console.log( + `[agentmemory] Loaded persisted vector index (${vectorIndex.size} vectors)`, + ); } const needsRebuild = - !loaded?.bm25 || - loaded.bm25.size === 0 || + bm25Index.size === 0 || (embeddingProvider && vectorIndex && vectorIndex.size === 0); if (needsRebuild) { diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 306633b..de00b0b 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -230,10 +230,11 @@ export function registerMcpEndpoints( sdk.registerFunction( { id: "mcp::tools::list" }, - async (): Promise => ({ - status_code: 200, - body: { tools: MCP_TOOLS }, - }), + async (req: ApiRequest): Promise => { + const authErr = checkAuth(req, secret); + if (authErr) return authErr; + return { status_code: 200, body: { tools: MCP_TOOLS } }; + }, ); sdk.registerTrigger({ type: "http", @@ -375,6 +376,7 @@ export function registerMcpEndpoints( ? (args.expandIds as string) .split(",") .map((id: string) => id.trim()) + .slice(0, 20) : []; const result = await sdk.trigger("mem::smart-search", { query: args.query, @@ -478,7 +480,7 @@ export function registerMcpEndpoints( return { status_code: 500, body: { - error: err instanceof Error ? err.message : "Internal error", + error: "Internal error", }, }; } diff --git a/src/state/search-index.ts b/src/state/search-index.ts index f9f08f8..9720a38 100644 --- a/src/state/search-index.ts +++ b/src/state/search-index.ts @@ -117,6 +117,13 @@ export class SearchIndex { this.totalDocLength = 0; } + restoreFrom(other: SearchIndex): void { + this.entries = other.entries; + this.invertedIndex = other.invertedIndex; + this.docTermCounts = other.docTermCounts; + this.totalDocLength = other.totalDocLength; + } + serialize(): string { const entries = Array.from(this.entries.entries()); const inverted = Array.from(this.invertedIndex.entries()).map( @@ -135,19 +142,24 @@ export class SearchIndex { } static deserialize(json: string): SearchIndex { - const idx = new SearchIndex(); - const data = JSON.parse(json); - for (const [key, val] of data.entries) { - idx.entries.set(key, val); - } - for (const [term, ids] of data.inverted) { - idx.invertedIndex.set(term, new Set(ids)); - } - for (const [id, counts] of data.docTerms) { - idx.docTermCounts.set(id, new Map(counts)); + try { + const idx = new SearchIndex(); + const data = JSON.parse(json); + if (!data?.entries || !data?.inverted || !data?.docTerms) return idx; + for (const [key, val] of data.entries) { + idx.entries.set(key, val); + } + for (const [term, ids] of data.inverted) { + idx.invertedIndex.set(term, new Set(ids)); + } + for (const [id, counts] of data.docTerms) { + idx.docTermCounts.set(id, new Map(counts)); + } + idx.totalDocLength = data.totalDocLength || 0; + return idx; + } catch { + return new SearchIndex(); } - idx.totalDocLength = data.totalDocLength; - return idx; } private extractTerms(obs: CompressedObservation): string[] { diff --git a/src/state/vector-index.ts b/src/state/vector-index.ts index 1bd0787..b24e362 100644 --- a/src/state/vector-index.ts +++ b/src/state/vector-index.ts @@ -21,10 +21,8 @@ function cosineSimilarity(a: Float32Array, b: Float32Array): number { } export class VectorIndex { - private vectors: Map< - string, - { embedding: Float32Array; sessionId: string } - > = new Map(); + private vectors: Map = + new Map(); add(obsId: string, sessionId: string, embedding: Float32Array): void { this.vectors.set(obsId, { embedding, sessionId }); @@ -61,10 +59,12 @@ export class VectorIndex { this.vectors.clear(); } + restoreFrom(other: VectorIndex): void { + this.vectors = (other as any).vectors; + } + serialize(): string { - const data: Array< - [string, { embedding: string; sessionId: string }] - > = []; + const data: Array<[string, { embedding: string; sessionId: string }]> = []; for (const [obsId, entry] of this.vectors) { data.push([ obsId, @@ -78,16 +78,20 @@ export class VectorIndex { } static deserialize(json: string): VectorIndex { - const idx = new VectorIndex(); - const data: Array< - [string, { embedding: string; sessionId: string }] - > = JSON.parse(json); - for (const [obsId, entry] of data) { - idx.vectors.set(obsId, { - embedding: base64ToFloat32(entry.embedding), - sessionId: entry.sessionId, - }); + try { + const idx = new VectorIndex(); + const data: Array<[string, { embedding: string; sessionId: string }]> = + JSON.parse(json); + if (!Array.isArray(data)) return idx; + for (const [obsId, entry] of data) { + idx.vectors.set(obsId, { + embedding: base64ToFloat32(entry.embedding), + sessionId: entry.sessionId, + }); + } + return idx; + } catch { + return new VectorIndex(); } - return idx; } } diff --git a/src/triggers/api.ts b/src/triggers/api.ts index 20078ee..cf36854 100644 --- a/src/triggers/api.ts +++ b/src/triggers/api.ts @@ -16,7 +16,7 @@ type Response = { }; const VIEWER_CSP = - "default-src 'self'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; connect-src 'self' ws://localhost:* wss://localhost:*"; + "default-src 'none'; script-src 'self'; style-src 'self'; connect-src 'self' ws://localhost:* wss://localhost:*; img-src 'self'; font-src 'self'"; function checkAuth( req: ApiRequest, @@ -37,27 +37,33 @@ export function registerApiTriggers( metricsStore?: MetricsStore, provider?: ResilientProvider | { circuitState?: unknown }, ): void { - sdk.registerFunction({ id: "api::health" }, async (): Promise => { - const health = await getLatestHealth(kv); - const functionMetrics = metricsStore ? await metricsStore.getAll() : []; - const circuitBreaker = - provider && "circuitState" in provider ? provider.circuitState : null; - - const status = health?.status || "healthy"; - const statusCode = status === "critical" ? 503 : 200; - - return { - status_code: statusCode, - body: { - status, - service: "agentmemory", - version: "0.3.0", - health: health || null, - functionMetrics, - circuitBreaker, - }, - }; - }); + sdk.registerFunction( + { id: "api::health" }, + async (req: ApiRequest): Promise => { + const authErr = checkAuth(req, secret); + if (authErr) return authErr; + + const health = await getLatestHealth(kv); + const functionMetrics = metricsStore ? await metricsStore.getAll() : []; + const circuitBreaker = + provider && "circuitState" in provider ? provider.circuitState : null; + + const status = health?.status || "healthy"; + const statusCode = status === "critical" ? 503 : 200; + + return { + status_code: statusCode, + body: { + status, + service: "agentmemory", + version: "0.3.0", + health: health || null, + functionMetrics, + circuitBreaker, + }, + }; + }, + ); sdk.registerTrigger({ type: "http", function_id: "api::health", @@ -183,10 +189,15 @@ export function registerApiTriggers( config: { api_path: "/agentmemory/summarize", http_method: "POST" }, }); - sdk.registerFunction({ id: "api::sessions" }, async (): Promise => { - const sessions = await kv.list(KV.sessions); - return { status_code: 200, body: { sessions } }; - }); + sdk.registerFunction( + { id: "api::sessions" }, + async (req: ApiRequest): Promise => { + const authErr = checkAuth(req, secret); + if (authErr) return authErr; + const sessions = await kv.list(KV.sessions); + return { status_code: 200, body: { sessions } }; + }, + ); sdk.registerTrigger({ type: "http", function_id: "api::sessions", @@ -196,6 +207,8 @@ export function registerApiTriggers( sdk.registerFunction( { id: "api::observations" }, async (req: ApiRequest): Promise => { + const authErr = checkAuth(req, secret); + if (authErr) return authErr; const sessionId = req.query_params["sessionId"] as string; if (!sessionId) return { status_code: 400, body: { error: "sessionId required" } }; From 7d29590b7b720871fba1e032edb68700aab10aac Mon Sep 17 00:00:00 2001 From: Rohit Ghumare Date: Fri, 27 Feb 2026 12:12:26 +0530 Subject: [PATCH 2/3] fix: address CodeRabbit PR #6 review findings - Deep copy in restoreFrom() to prevent shared mutable state - Shape validation for import payload (Array.isArray checks) - Global observation cap (500K) on import - safeParseInt() using Number.isNaN() to preserve explicit zero - Unauthenticated /agentmemory/livez liveness endpoint - CSP reverted to unsafe-inline for viewer inline code - Sort memories by createdAt before slicing in auto-forget - Truncation signal in smart-search expanded results - Removed false vector index trigger from needsRebuild --- src/config.ts | 11 +++++++--- src/functions/auto-forget.ts | 4 ++++ src/functions/export-import.ts | 37 ++++++++++++++++++++++++++++++---- src/functions/smart-search.ts | 7 +++++-- src/index.ts | 4 +--- src/state/search-index.ts | 16 ++++++++++++--- src/state/vector-index.ts | 2 +- src/triggers/api.ts | 15 +++++++++++++- 8 files changed, 79 insertions(+), 17 deletions(-) diff --git a/src/config.ts b/src/config.ts index 2858142..bf80a0a 100644 --- a/src/config.ts +++ b/src/config.ts @@ -8,6 +8,12 @@ import type { FallbackConfig, } from "./types.js"; +function safeParseInt(value: string | undefined, fallback: number): number { + if (!value) return fallback; + const parsed = parseInt(value, 10); + return Number.isNaN(parsed) ? fallback : parsed; +} + const DATA_DIR = join(homedir(), ".agentmemory"); const ENV_FILE = join(DATA_DIR, ".env"); @@ -74,9 +80,8 @@ export function loadConfig(): AgentMemoryConfig { restPort: parseInt(env["III_REST_PORT"] || "3111", 10) || 3111, streamsPort: parseInt(env["III_STREAMS_PORT"] || "3112", 10) || 3112, provider, - tokenBudget: parseInt(env["TOKEN_BUDGET"] || "2000", 10) || 2000, - maxObservationsPerSession: - parseInt(env["MAX_OBS_PER_SESSION"] || "500", 10) || 500, + tokenBudget: safeParseInt(env["TOKEN_BUDGET"], 2000), + maxObservationsPerSession: safeParseInt(env["MAX_OBS_PER_SESSION"], 500), compressionModel: provider.model, dataDir: DATA_DIR, }; diff --git a/src/functions/auto-forget.ts b/src/functions/auto-forget.ts index 96ba575..ce90d83 100644 --- a/src/functions/auto-forget.ts +++ b/src/functions/auto-forget.ts @@ -54,6 +54,10 @@ export function registerAutoForgetFunction(sdk: ISdk, kv: StateKV): void { const latestMemories = memories .filter((m) => m.isLatest !== false && !deletedIds.has(m.id)) + .sort( + (a, b) => + new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime(), + ) .slice(0, 1000); for (let i = 0; i < latestMemories.length; i++) { for (let j = i + 1; j < latestMemories.length; j++) { diff --git a/src/functions/export-import.ts b/src/functions/export-import.ts index 0fa0788..48406a8 100644 --- a/src/functions/export-import.ts +++ b/src/functions/export-import.ts @@ -89,32 +89,61 @@ export function registerExportImportFunction(sdk: ISdk, kv: StateKV): void { const MAX_MEMORIES = 50_000; const MAX_SUMMARIES = 10_000; const MAX_OBS_PER_SESSION = 5_000; + const MAX_TOTAL_OBSERVATIONS = 500_000; - if (importData.sessions?.length > MAX_SESSIONS) { + if (!Array.isArray(importData.sessions)) { + return { success: false, error: "sessions must be an array" }; + } + if (!Array.isArray(importData.memories)) { + return { success: false, error: "memories must be an array" }; + } + if (!Array.isArray(importData.summaries)) { + return { success: false, error: "summaries must be an array" }; + } + if ( + typeof importData.observations !== "object" || + importData.observations === null || + Array.isArray(importData.observations) + ) { + return { success: false, error: "observations must be an object" }; + } + + if (importData.sessions.length > MAX_SESSIONS) { return { success: false, error: `Too many sessions (max ${MAX_SESSIONS})`, }; } - if (importData.memories?.length > MAX_MEMORIES) { + if (importData.memories.length > MAX_MEMORIES) { return { success: false, error: `Too many memories (max ${MAX_MEMORIES})`, }; } - if (importData.summaries?.length > MAX_SUMMARIES) { + if (importData.summaries.length > MAX_SUMMARIES) { return { success: false, error: `Too many summaries (max ${MAX_SUMMARIES})`, }; } - for (const [, obs] of Object.entries(importData.observations || {})) { + let totalObservations = 0; + for (const [, obs] of Object.entries(importData.observations)) { + if (!Array.isArray(obs)) { + return { success: false, error: "observation values must be arrays" }; + } if (obs.length > MAX_OBS_PER_SESSION) { return { success: false, error: `Too many observations per session (max ${MAX_OBS_PER_SESSION})`, }; } + totalObservations += obs.length; + } + if (totalObservations > MAX_TOTAL_OBSERVATIONS) { + return { + success: false, + error: `Too many total observations (max ${MAX_TOTAL_OBSERVATIONS})`, + }; } const stats = { diff --git a/src/functions/smart-search.ts b/src/functions/smart-search.ts index 79c59d4..0e7eec4 100644 --- a/src/functions/smart-search.ts +++ b/src/functions/smart-search.ts @@ -41,11 +41,14 @@ export function registerSmartSearchFunction( } } + const truncated = data.expandIds.length > ids.length; ctx.logger.info("Smart search expanded", { - requested: ids.length, + requested: data.expandIds.length, + returned: ids.length, found: expanded.length, + truncated, }); - return { mode: "expanded", results: expanded }; + return { mode: "expanded", results: expanded, truncated }; } if (!data.query || typeof data.query !== "string" || !data.query.trim()) { diff --git a/src/index.ts b/src/index.ts index 4bc6c5f..ac3f6b7 100644 --- a/src/index.ts +++ b/src/index.ts @@ -153,9 +153,7 @@ async function main() { ); } - const needsRebuild = - bm25Index.size === 0 || - (embeddingProvider && vectorIndex && vectorIndex.size === 0); + const needsRebuild = bm25Index.size === 0; if (needsRebuild) { const indexCount = await rebuildIndex(kv).catch((err) => { diff --git a/src/state/search-index.ts b/src/state/search-index.ts index 9720a38..e937e4e 100644 --- a/src/state/search-index.ts +++ b/src/state/search-index.ts @@ -118,9 +118,19 @@ export class SearchIndex { } restoreFrom(other: SearchIndex): void { - this.entries = other.entries; - this.invertedIndex = other.invertedIndex; - this.docTermCounts = other.docTermCounts; + this.entries = new Map(other.entries); + this.invertedIndex = new Map( + Array.from(other.invertedIndex.entries()).map(([k, v]) => [ + k, + new Set(v), + ]), + ); + this.docTermCounts = new Map( + Array.from(other.docTermCounts.entries()).map(([k, v]) => [ + k, + new Map(v), + ]), + ); this.totalDocLength = other.totalDocLength; } diff --git a/src/state/vector-index.ts b/src/state/vector-index.ts index b24e362..6624770 100644 --- a/src/state/vector-index.ts +++ b/src/state/vector-index.ts @@ -60,7 +60,7 @@ export class VectorIndex { } restoreFrom(other: VectorIndex): void { - this.vectors = (other as any).vectors; + this.vectors = new Map((other as any).vectors); } serialize(): string { diff --git a/src/triggers/api.ts b/src/triggers/api.ts index cf36854..3d555d8 100644 --- a/src/triggers/api.ts +++ b/src/triggers/api.ts @@ -16,7 +16,7 @@ type Response = { }; const VIEWER_CSP = - "default-src 'none'; script-src 'self'; style-src 'self'; connect-src 'self' ws://localhost:* wss://localhost:*; img-src 'self'; font-src 'self'"; + "default-src 'none'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; connect-src 'self' ws://localhost:* wss://localhost:*; img-src 'self'; font-src 'self'"; function checkAuth( req: ApiRequest, @@ -37,6 +37,19 @@ export function registerApiTriggers( metricsStore?: MetricsStore, provider?: ResilientProvider | { circuitState?: unknown }, ): void { + sdk.registerFunction( + { id: "api::liveness" }, + async (): Promise => ({ + status_code: 200, + body: { status: "ok", service: "agentmemory" }, + }), + ); + sdk.registerTrigger({ + type: "http", + function_id: "api::liveness", + config: { api_path: "/agentmemory/livez", http_method: "GET" }, + }); + sdk.registerFunction( { id: "api::health" }, async (req: ApiRequest): Promise => { From 8c66484c9be8aedc656cf7ef043d18c52ca7ef98 Mon Sep 17 00:00:00 2001 From: Rohit Ghumare Date: Fri, 27 Feb 2026 12:23:31 +0530 Subject: [PATCH 3/3] fix: address CodeRabbit PR #6 round 3 findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add MAX_OBS_BUCKETS (10K) cap on import observation keys - Fix smart-search log: returned → expanded.length (actual response size) - Validate totalDocLength in deserialize (finite, non-negative, floor) - Deep copy Float32Array in VectorIndex.restoreFrom() - Clone IndexEntry objects in SearchIndex.restoreFrom() - Per-row resilient VectorIndex.deserialize() (skip bad rows, keep good) --- src/functions/export-import.ts | 9 ++++++++ src/functions/smart-search.ts | 4 ++-- src/state/search-index.ts | 8 +++++-- src/state/vector-index.ts | 40 ++++++++++++++++++++++++++-------- 4 files changed, 48 insertions(+), 13 deletions(-) diff --git a/src/functions/export-import.ts b/src/functions/export-import.ts index 48406a8..ac9dc41 100644 --- a/src/functions/export-import.ts +++ b/src/functions/export-import.ts @@ -126,6 +126,15 @@ export function registerExportImportFunction(sdk: ISdk, kv: StateKV): void { error: `Too many summaries (max ${MAX_SUMMARIES})`, }; } + const MAX_OBS_BUCKETS = 10_000; + const obsBuckets = Object.keys(importData.observations); + if (obsBuckets.length > MAX_OBS_BUCKETS) { + return { + success: false, + error: `Too many observation buckets (max ${MAX_OBS_BUCKETS})`, + }; + } + let totalObservations = 0; for (const [, obs] of Object.entries(importData.observations)) { if (!Array.isArray(obs)) { diff --git a/src/functions/smart-search.ts b/src/functions/smart-search.ts index 0e7eec4..d525d6a 100644 --- a/src/functions/smart-search.ts +++ b/src/functions/smart-search.ts @@ -44,8 +44,8 @@ export function registerSmartSearchFunction( const truncated = data.expandIds.length > ids.length; ctx.logger.info("Smart search expanded", { requested: data.expandIds.length, - returned: ids.length, - found: expanded.length, + attempted: ids.length, + returned: expanded.length, truncated, }); return { mode: "expanded", results: expanded, truncated }; diff --git a/src/state/search-index.ts b/src/state/search-index.ts index e937e4e..2427b58 100644 --- a/src/state/search-index.ts +++ b/src/state/search-index.ts @@ -118,7 +118,9 @@ export class SearchIndex { } restoreFrom(other: SearchIndex): void { - this.entries = new Map(other.entries); + this.entries = new Map( + Array.from(other.entries.entries()).map(([k, v]) => [k, { ...v }]), + ); this.invertedIndex = new Map( Array.from(other.invertedIndex.entries()).map(([k, v]) => [ k, @@ -165,7 +167,9 @@ export class SearchIndex { for (const [id, counts] of data.docTerms) { idx.docTermCounts.set(id, new Map(counts)); } - idx.totalDocLength = data.totalDocLength || 0; + const rawLen = Number(data.totalDocLength); + idx.totalDocLength = + Number.isFinite(rawLen) && rawLen >= 0 ? Math.floor(rawLen) : 0; return idx; } catch { return new SearchIndex(); diff --git a/src/state/vector-index.ts b/src/state/vector-index.ts index 6624770..7a8f442 100644 --- a/src/state/vector-index.ts +++ b/src/state/vector-index.ts @@ -60,7 +60,17 @@ export class VectorIndex { } restoreFrom(other: VectorIndex): void { - this.vectors = new Map((other as any).vectors); + const src = (other as any).vectors as Map< + string, + { embedding: Float32Array; sessionId: string } + >; + this.vectors = new Map(); + for (const [obsId, entry] of src) { + this.vectors.set(obsId, { + embedding: new Float32Array(entry.embedding), + sessionId: entry.sessionId, + }); + } } serialize(): string { @@ -78,20 +88,32 @@ export class VectorIndex { } static deserialize(json: string): VectorIndex { + const idx = new VectorIndex(); + let data: unknown; try { - const idx = new VectorIndex(); - const data: Array<[string, { embedding: string; sessionId: string }]> = - JSON.parse(json); - if (!Array.isArray(data)) return idx; - for (const [obsId, entry] of data) { + data = JSON.parse(json); + } catch { + return idx; + } + if (!Array.isArray(data)) return idx; + for (const row of data) { + try { + if (!Array.isArray(row) || row.length < 2) continue; + const [obsId, entry] = row; + if ( + typeof obsId !== "string" || + typeof entry?.embedding !== "string" || + typeof entry?.sessionId !== "string" + ) + continue; idx.vectors.set(obsId, { embedding: base64ToFloat32(entry.embedding), sessionId: entry.sessionId, }); + } catch { + continue; } - return idx; - } catch { - return new VectorIndex(); } + return idx; } }