diff --git a/docs/plugins/sdk-overview.md b/docs/plugins/sdk-overview.md index 5c69e00ce77..e3f33cb087f 100644 --- a/docs/plugins/sdk-overview.md +++ b/docs/plugins/sdk-overview.md @@ -113,6 +113,16 @@ can consume this generic provider surface. The older `contracts.memoryEmbeddingProviders` seam is deprecated compatibility while existing memory-specific providers migrate. +Memory-specific providers that still expose a runtime `batchEmbed(...)` stay on +the existing per-file batching contract unless their runtime explicitly sets +`sourceWideBatchEmbed: true`. That opt-in lets the memory host submit chunks from +multiple dirty memory files and enabled sources in one `batchEmbed(...)` call up +to the host batch limits. Batch adapters that upload JSONL request files must +split provider jobs before their upload-size cap as well as their request-count +cap. The provider must return one embedding per input chunk in the same order as +`batch.chunks`; omit the flag when the provider expects file-local batches or +cannot preserve input ordering across a larger source-wide job. + ### Tools and commands Use [`defineToolPlugin`](/plugins/tool-plugins) for simple tool-only plugins diff --git a/extensions/memory-core/src/memory/index.test.ts b/extensions/memory-core/src/memory/index.test.ts index 495527068df..7d6c9170d3e 100644 --- a/extensions/memory-core/src/memory/index.test.ts +++ b/extensions/memory-core/src/memory/index.test.ts @@ -13,6 +13,7 @@ import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } import "./test-runtime-mocks.js"; import type { MemoryIndexManager } from "./index.js"; import { closeAllMemorySearchManagers, getMemorySearchManager } from "./index.js"; +import { splitSourceWideEmbeddingChunks } from "./manager-embedding-ops.js"; import { LOCAL_EMBEDDING_WORKER_ERROR_CODES } from "./manager-local-worker-errors.js"; import type { MemoryIndexMeta } from "./manager-reindex-state.js"; import { closeMemoryIndexManagersForAgent, EMBEDDING_PROBE_CACHE_TTL_MS } from "./manager.js"; @@ -28,6 +29,11 @@ afterAll(() => { let embedBatchCalls = 0; let embedBatchInputCalls = 0; +let providerRuntimeBatchCalls: string[][] = []; +let providerRuntimeBatchGate: Promise | null = null; +let providerRuntimeBatchFailuresRemaining = 0; +let providerRuntimeActiveBatchCalls = 0; +let providerRuntimeMaxActiveBatchCalls = 0; let providerCloseCalls = 0; let providerCloseFailuresRemaining = 0; let providerCloseGate: Promise | null = null; @@ -88,6 +94,8 @@ vi.mock("./embeddings.js", () => { const providerId = options.provider === "gemini" || options.provider === "fallback-provider" || + options.provider === "batch-test" || + options.provider === "batch-wide-test" || options.provider === "ollama" ? options.provider : "mock"; @@ -141,20 +149,45 @@ vi.mock("./embeddings.js", () => { } : {}), }, - ...(providerId === "gemini" || providerId === "fallback-provider" + ...(providerId === "batch-test" || providerId === "batch-wide-test" ? { runtime: { id: providerId, - cacheKeyData: { - provider: providerId, - baseUrl: "https://generativelanguage.googleapis.com/v1beta", - model, - outputDimensionality: options.outputDimensionality, - headers: [], + ...(providerId === "batch-wide-test" ? { sourceWideBatchEmbed: true } : {}), + batchEmbed: async (batch: { chunks: Array<{ text: string }> }) => { + providerRuntimeActiveBatchCalls += 1; + providerRuntimeMaxActiveBatchCalls = Math.max( + providerRuntimeMaxActiveBatchCalls, + providerRuntimeActiveBatchCalls, + ); + try { + await providerRuntimeBatchGate; + providerRuntimeBatchCalls.push(batch.chunks.map((chunk) => chunk.text)); + if (providerRuntimeBatchFailuresRemaining > 0) { + providerRuntimeBatchFailuresRemaining -= 1; + throw new Error("provider runtime batch failed"); + } + return batch.chunks.map((chunk) => embedText(chunk.text)); + } finally { + providerRuntimeActiveBatchCalls -= 1; + } }, }, } - : {}), + : providerId === "gemini" || providerId === "fallback-provider" + ? { + runtime: { + id: providerId, + cacheKeyData: { + provider: providerId, + baseUrl: "https://generativelanguage.googleapis.com/v1beta", + model, + outputDimensionality: options.outputDimensionality, + headers: [], + }, + }, + } + : {}), }; }, }; @@ -221,6 +254,11 @@ describe("memory index", () => { registerBuiltInMemoryEmbeddingProviders({ registerMemoryEmbeddingProvider: registerAdapter }); embedBatchCalls = 0; embedBatchInputCalls = 0; + providerRuntimeBatchCalls = []; + providerRuntimeBatchGate = null; + providerRuntimeBatchFailuresRemaining = 0; + providerRuntimeActiveBatchCalls = 0; + providerRuntimeMaxActiveBatchCalls = 0; providerCloseCalls = 0; providerCloseFailuresRemaining = 0; providerCloseGate = null; @@ -268,6 +306,7 @@ describe("memory index", () => { provider?: string; fallback?: "none" | "gemini" | "fallback-provider"; providerAliases?: NonNullable["providers"]>; + batchEnabled?: boolean; model?: string; outputDimensionality?: number; multimodal?: { @@ -294,6 +333,12 @@ describe("memory index", () => { // Perf: keep test indexes to a single chunk to reduce sqlite work. chunking: { tokens: 4000, overlap: 0 }, sync: { watch: false, onSessionStart: false, onSearch: params.onSearch ?? true }, + remote: params.batchEnabled + ? { + nonBatchConcurrency: 1, + batch: { enabled: true, pollIntervalMs: 0, timeoutMinutes: 1 }, + } + : undefined, query: { minScore: params.minScore ?? 0, hybrid: params.hybrid ?? { enabled: false }, @@ -419,6 +464,243 @@ describe("memory index", () => { } }); + it("batches dirty memory chunks across files", async () => { + await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line."); + await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nGamma memory line."); + const cfg = createCfg({ + provider: "batch-wide-test", + batchEnabled: true, + storePath: path.join(workspaceDir, "index-cross-file-batch.sqlite"), + }); + const manager = await getFreshManager(cfg); + try { + await manager.sync({ reason: "test" }); + + expect(providerRuntimeBatchCalls).toHaveLength(1); + expect(providerRuntimeBatchCalls[0]).toEqual([ + "# Log\nAlpha memory line.\nZebra memory line.", + "# Log\nBeta memory line.", + "# Log\nGamma memory line.", + ]); + } finally { + await manager.close?.(); + } + }); + + it("maps source-wide batch fallback results to missing chunks after cache hits", async () => { + const cfg = createCfg({ + provider: "batch-wide-test", + batchEnabled: true, + storePath: path.join(workspaceDir, "index-cross-file-batch-fallback-cache.sqlite"), + }); + const manager = await getFreshManager(cfg); + try { + await manager.sync({ reason: "test" }); + + await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line."); + providerRuntimeBatchCalls = []; + providerRuntimeBatchFailuresRemaining = 1; + embedBatchCalls = 0; + + await manager.sync({ reason: "test", force: true }); + + expect(providerRuntimeBatchCalls).toEqual([["# Log\nBeta memory line."]]); + expect(embedBatchCalls).toBe(1); + const betaRow = ( + manager as unknown as { + db: { prepare: (sql: string) => { get: (...args: unknown[]) => unknown } }; + } + ).db + .prepare("SELECT embedding FROM chunks WHERE path LIKE ? AND source = ?") + .get("%2026-01-13.md", "memory") as { embedding: string } | undefined; + + expect(betaRow).toBeDefined(); + expect(JSON.parse(betaRow?.embedding ?? "[]")).toEqual([0, 1, 0, 0]); + } finally { + await manager.close?.(); + } + }); + + it("splits oversized source-wide embedding requests at the request cap", () => { + expect(splitSourceWideEmbeddingChunks(["one", "two", "three", "four", "five"], 2)).toEqual([ + ["one", "two"], + ["three", "four"], + ["five"], + ]); + }); + + it("keeps split chunks from oversized files in one source-wide batch", async () => { + await fs.writeFile( + path.join(memoryDir, "2026-01-13.md"), + `# Log\n${"Long split memory line. ".repeat(1200)}`, + ); + await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nBeta memory line."); + const cfg = createCfg({ + provider: "batch-wide-test", + batchEnabled: true, + storePath: path.join(workspaceDir, "index-split-chunks-cross-file-batch.sqlite"), + }); + const manager = await getFreshManager(cfg); + try { + await manager.sync({ reason: "test" }); + + expect(providerRuntimeBatchCalls).toHaveLength(1); + const combinedBatch = providerRuntimeBatchCalls[0] ?? []; + expect(combinedBatch.length).toBeGreaterThan(3); + expect(combinedBatch.join("\n")).toContain("Long split memory line."); + expect(combinedBatch).toContain("# Log\nBeta memory line."); + } finally { + await manager.close?.(); + } + }); + + it("keeps custom batch runtimes per file without source-wide opt in", async () => { + await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line."); + await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nGamma memory line."); + const cfg = createCfg({ + provider: "batch-test", + batchEnabled: true, + storePath: path.join(workspaceDir, "index-custom-batch-compat.sqlite"), + }); + const manager = await getFreshManager(cfg); + try { + await manager.sync({ reason: "test" }); + + expect(providerRuntimeBatchCalls).toHaveLength(3); + expect(providerRuntimeBatchCalls.every((call) => call.length === 1)).toBe(true); + expect(providerRuntimeBatchCalls.map((call) => call[0]).toSorted()).toEqual( + [ + "# Log\nAlpha memory line.\nZebra memory line.", + "# Log\nBeta memory line.", + "# Log\nGamma memory line.", + ].toSorted(), + ); + } finally { + await manager.close?.(); + } + }); + + it("keeps custom batch runtimes concurrent without source-wide opt in", async () => { + await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line."); + await fs.writeFile(path.join(memoryDir, "2026-01-14.md"), "# Log\nGamma memory line."); + const cfg = createCfg({ + provider: "batch-test", + batchEnabled: true, + storePath: path.join(workspaceDir, "index-custom-batch-concurrency.sqlite"), + }); + const manager = await getFreshManager(cfg); + let releaseBatchGate: (() => void) | undefined; + providerRuntimeBatchGate = new Promise((resolve) => { + releaseBatchGate = resolve; + }); + const syncPromise = manager.sync({ reason: "test" }); + let waitError: Error | undefined; + try { + await vi.waitFor(() => expect(providerRuntimeMaxActiveBatchCalls).toBeGreaterThan(1)); + } catch (err) { + waitError = err instanceof Error ? err : new Error(String(err)); + } finally { + releaseBatchGate?.(); + await syncPromise; + await manager.close?.(); + } + if (waitError) { + throw waitError; + } + }); + + it("bounds source-wide memory batches", async () => { + const batchFileLimit = 2048; + for (let index = 0; index < batchFileLimit; index += 1) { + await fs.writeFile( + path.join(memoryDir, `2026-02-${String(index + 1).padStart(4, "0")}.md`), + `# Log\nBounded memory line ${index}.`, + ); + } + const cfg = createCfg({ + provider: "batch-wide-test", + batchEnabled: true, + storePath: path.join(workspaceDir, "index-bounded-cross-file-batch.sqlite"), + }); + const manager = await getFreshManager(cfg); + try { + await manager.sync({ reason: "test" }); + + expect(providerRuntimeBatchCalls).toHaveLength(2); + expect(providerRuntimeBatchCalls[0]).toHaveLength(batchFileLimit); + expect(providerRuntimeBatchCalls[1]).toHaveLength(1); + expect(providerRuntimeBatchCalls.flat()).toHaveLength(batchFileLimit + 1); + } finally { + await manager.close?.(); + } + }); + + it("batches forced memory and session indexing across files", async () => { + await fs.writeFile(path.join(memoryDir, "2026-01-13.md"), "# Log\nBeta memory line."); + const sessionsDir = resolveSessionTranscriptsDirForAgent("main"); + await fs.mkdir(sessionsDir, { recursive: true }); + await fs.writeFile( + path.join(sessionsDir, "session-alpha.jsonl"), + [ + JSON.stringify({ + type: "session", + id: "session-alpha", + timestamp: "2026-04-07T15:24:04.113Z", + }), + JSON.stringify({ + type: "message", + message: { + role: "user", + timestamp: "2026-04-07T15:25:04.113Z", + content: [{ type: "text", text: "Session alpha memory line." }], + }, + }), + ].join("\n") + "\n", + "utf8", + ); + await fs.writeFile( + path.join(sessionsDir, "session-beta.jsonl"), + [ + JSON.stringify({ + type: "session", + id: "session-beta", + timestamp: "2026-04-07T15:24:04.113Z", + }), + JSON.stringify({ + type: "message", + message: { + role: "assistant", + timestamp: "2026-04-07T15:25:04.113Z", + content: [{ type: "text", text: "Session beta memory line." }], + }, + }), + ].join("\n") + "\n", + "utf8", + ); + const cfg = createCfg({ + provider: "batch-wide-test", + batchEnabled: true, + sources: ["memory", "sessions"], + sessionMemory: true, + storePath: path.join(workspaceDir, "index-force-cross-source-batch.sqlite"), + }); + const manager = await getFreshManager(cfg); + try { + await manager.sync({ reason: "cli", force: true }); + + expect(providerRuntimeBatchCalls).toHaveLength(1); + const combinedBatch = providerRuntimeBatchCalls[0] ?? []; + expect(combinedBatch.slice(0, 2)).toEqual([ + "# Log\nAlpha memory line.\nZebra memory line.", + "# Log\nBeta memory line.", + ]); + expect(combinedBatch.join("\n")).toContain("Session alpha memory line."); + expect(combinedBatch.join("\n")).toContain("Session beta memory line."); + } finally { + await manager.close?.(); + } + }); + it("does not full-reindex on search when existing metadata belongs to another provider", async () => { const dbPath = path.join(workspaceDir, "index-provider-cutover.sqlite"); const oldCfg = createCfg({ @@ -1041,25 +1323,6 @@ describe("memory index", () => { expect(third).toBe(second); }); - it("closes stale default managers when provider requirement changes", async () => { - const storePath = path.join(workspaceDir, "index-provider-requirement-cache.sqlite"); - const implicitCfg = createCfg({ storePath }); - const implicit = requireManager( - await getMemorySearchManager({ cfg: implicitCfg, agentId: "main" }), - ); - managersForCleanup.add(implicit); - await implicit.probeEmbeddingAvailability(); - - const explicitCfg = createCfg({ storePath, provider: "openai" }); - const explicit = requireManager( - await getMemorySearchManager({ cfg: explicitCfg, agentId: "main" }), - ); - managersForCleanup.add(explicit); - - expect(explicit === implicit).toBe(false); - expect(providerCloseCalls).toBe(1); - }); - it("retries embedding provider close before releasing the manager", async () => { providerCloseFailuresRemaining = 1; const cfg = createCfg({ @@ -1599,6 +1862,25 @@ describe("memory index", () => { } }); + it("keeps metadata after unchanged safe force reindex", async () => { + vi.stubEnv("OPENCLAW_TEST_MEMORY_UNSAFE_REINDEX", "0"); + const cfg = createCfg({ + storePath: path.join(workspaceDir, "index-safe-force-metadata.sqlite"), + }); + const manager = await getFreshManager(cfg); + try { + await manager.sync({ reason: "test", force: true }); + expect(manager.status().custom?.indexIdentity).toEqual({ status: "valid" }); + + await manager.sync({ reason: "cli", force: true }); + + expect(manager.status().dirty).toBe(false); + expect(manager.status().custom?.indexIdentity).toEqual({ status: "valid" }); + } finally { + await manager.close?.(); + } + }); + it("streams embedding cache rows during safe reindex", async () => { vi.stubEnv("OPENCLAW_TEST_MEMORY_UNSAFE_REINDEX", "0"); type EmbeddingCacheRow = { @@ -1690,12 +1972,6 @@ describe("memory index", () => { const status = manager.status(); expect(status.chunks).toBeGreaterThan(0); expect(embedBatchCalls).toBe(0); - expect(status.custom?.providerUnavailableReason).toBe("No API key found for provider"); - expect(status.custom?.providerState).toEqual({ - mode: "fts-only", - reason: "No API key found for provider", - attemptedProviderId: "openai", - }); const results = await manager.search("Alpha"); expect(results.length).toBeGreaterThan(0); @@ -1754,7 +2030,6 @@ describe("memory index", () => { await manager.close?.(); } }); - it("prefers exact session transcript hits in FTS-only mode", async () => { try { const manager = await getFtsSessionManager({ diff --git a/extensions/memory-core/src/memory/manager-embedding-ops.ts b/extensions/memory-core/src/memory/manager-embedding-ops.ts index 46f1b92f312..01e84fbbb41 100644 --- a/extensions/memory-core/src/memory/manager-embedding-ops.ts +++ b/extensions/memory-core/src/memory/manager-embedding-ops.ts @@ -14,6 +14,7 @@ import { hashText, remapChunkLines, retryTransientMemoryRead, + runWithConcurrency, type MemoryChunk, type MemorySource, } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; @@ -40,7 +41,7 @@ import { runMemoryEmbeddingRetryLoop, } from "./manager-embedding-policy.js"; import { deleteMemoryFtsRows } from "./manager-fts-state.js"; -import { MemoryManagerSyncOps } from "./manager-sync-ops.js"; +import { MemoryManagerSyncOps, type MemoryIndexWorkItem } from "./manager-sync-ops.js"; import { logMemoryVectorDegradedWrite } from "./manager-vector-warning.js"; import { replaceMemoryVectorRow } from "./manager-vector-write.js"; @@ -56,6 +57,8 @@ const EMBEDDING_QUERY_TIMEOUT_REMOTE_MS = 60_000; const EMBEDDING_QUERY_TIMEOUT_LOCAL_MS = 5 * 60_000; const EMBEDDING_BATCH_TIMEOUT_REMOTE_MS = 2 * 60_000; const EMBEDDING_BATCH_TIMEOUT_LOCAL_MS = 10 * 60_000; +const SOURCE_WIDE_BATCH_MAX_FILES = 2048; +const SOURCE_WIDE_BATCH_MAX_REQUESTS = 50000; const log = createSubsystemLogger("memory"); @@ -70,17 +73,46 @@ function resolveEmbeddingSecondsTimeoutMs(seconds: number): number { ); } -type MemoryIndexEntry = { - path: string; - absPath: string; - mtimeMs: number; - size: number; - hash: string; - kind?: "markdown" | "multimodal"; - contentText?: string; - lineMap?: number[]; +type MemoryIndexEntry = MemoryIndexWorkItem["entry"]; + +type PreparedMemoryIndexEntry = { + entry: MemoryIndexEntry; + source: MemorySource; + chunks: MemoryChunk[]; + structuredInputBytes?: number; }; +function countBatchSources(items: Array<{ source: MemorySource }>): Record { + const counts: Record = {}; + for (const item of items) { + counts[item.source] = (counts[item.source] ?? 0) + 1; + } + return counts; +} + +function formatBatchSourceLabel(counts: Record): string { + const sources = Object.keys(counts).toSorted(); + return sources.length > 0 ? sources.join("+") : "unknown"; +} + +function formatBatchSourceCounts(counts: Record): string { + return ( + Object.entries(counts) + .toSorted(([left], [right]) => left.localeCompare(right)) + .map(([source, count]) => `${source}=${count}`) + .join(",") || "none" + ); +} + +export function splitSourceWideEmbeddingChunks(chunks: T[], maxRequests: number): T[][] { + const limit = Math.max(1, Math.floor(maxRequests)); + const batches: T[][] = []; + for (let start = 0; start < chunks.length; start += limit) { + batches.push(chunks.slice(start, start + limit)); + } + return batches; +} + export function resolveEmbeddingTimeoutMs(params: { kind: "query" | "batch"; providerId?: string; @@ -256,18 +288,25 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { return hashText(JSON.stringify({ provider: this.provider.id, model: this.provider.model })); } - private buildBatchDebug(source: MemorySource, chunks: MemoryChunk[]) { + private buildBatchDebug( + source: string, + chunks: MemoryChunk[], + context: Record = {}, + ) { return (message: string, data?: Record) => log.debug( message, - data ? { ...data, source, chunks: chunks.length } : { source, chunks: chunks.length }, + data + ? { ...data, source, chunks: chunks.length, ...context } + : { source, chunks: chunks.length, ...context }, ); } private async embedChunksWithBatch( chunks: MemoryChunk[], _entry: MemoryIndexEntry, - source: MemorySource, + source: string, + debugContext: Record = {}, ): Promise { const provider = this.provider; const batchEmbed = this.providerRuntime?.batchEmbed; @@ -293,9 +332,9 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { concurrency: this.batch.concurrency, pollIntervalMs: this.batch.pollIntervalMs, timeoutMs: this.batch.timeoutMs, - debug: this.buildBatchDebug(source, chunks), + debug: this.buildBatchDebug(source, chunks, debugContext), }), - fallback: async () => await this.embedChunksInBatches(chunks), + fallback: async () => await this.embedChunksInBatches(missingChunks), }); if (!batchResult) { return this.embedChunksInBatches(chunks); @@ -747,6 +786,165 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { this.upsertFileRecord(entry, source); } + private async prepareIndexEntry( + entry: MemoryIndexEntry, + options: { source: MemorySource; content?: string }, + ): Promise { + if ("kind" in entry && entry.kind === "multimodal") { + const multimodalChunk = await buildMultimodalChunkForIndexing(entry); + if (!multimodalChunk) { + this.clearIndexedFileData(entry.path, options.source); + this.deleteFileRecord(entry.path, options.source); + return null; + } + return { + entry, + source: options.source, + chunks: [multimodalChunk.chunk], + structuredInputBytes: multimodalChunk.structuredInputBytes, + }; + } + + const content = + options.content ?? + entry.content ?? + (await retryTransientMemoryRead( + () => fs.readFile(entry.absPath, "utf-8"), + `read memory markdown for indexing ${entry.absPath}`, + )); + const baseChunks = filterNonEmptyMemoryChunks(chunkMarkdown(content, this.settings.chunking)); + const chunks = this.provider + ? enforceEmbeddingMaxInputTokens(this.provider, baseChunks, EMBEDDING_BATCH_MAX_TOKENS) + : baseChunks; + if (options.source === "sessions" && "lineMap" in entry) { + remapChunkLines(chunks, entry.lineMap); + } + return { entry, source: options.source, chunks }; + } + + protected override async indexFiles(items: MemoryIndexWorkItem[]): Promise { + if (items.length === 0) { + return; + } + const provider = this.provider; + const batchEmbed = this.providerRuntime?.batchEmbed; + if ( + !provider || + !this.batch.enabled || + !batchEmbed || + this.providerRuntime?.sourceWideBatchEmbed !== true + ) { + await runWithConcurrency( + items.map((item) => async () => await this.indexFile(item.entry, { source: item.source })), + this.getIndexConcurrency(), + ); + return; + } + + const itemSourceCounts = countBatchSources(items); + log.debug( + `memory embeddings: source-wide batch prepare files=${items.length} sources=${formatBatchSourceCounts( + itemSourceCounts, + )} maxFiles=${SOURCE_WIDE_BATCH_MAX_FILES} maxRequests=${SOURCE_WIDE_BATCH_MAX_REQUESTS}`, + { + files: items.length, + sources: itemSourceCounts, + maxFiles: SOURCE_WIDE_BATCH_MAX_FILES, + maxRequests: SOURCE_WIDE_BATCH_MAX_REQUESTS, + }, + ); + + let prepared: PreparedMemoryIndexEntry[] = []; + let preparedRequestCount = 0; + let sourceWideBatchGroup = 0; + const flushPrepared = async (reason: "max-files" | "max-requests" | "end") => { + const firstEntry = prepared[0]?.entry; + if (!firstEntry) { + return; + } + const current = prepared; + const chunks = current.flatMap((item) => item.chunks); + const sourceCounts = countBatchSources(current); + const source = formatBatchSourceLabel(sourceCounts); + sourceWideBatchGroup += 1; + const chunkBatches = splitSourceWideEmbeddingChunks(chunks, SOURCE_WIDE_BATCH_MAX_REQUESTS); + log.debug( + `memory embeddings: source-wide batch submit group=${sourceWideBatchGroup} source=${source} files=${current.length} chunks=${chunks.length} requests=${chunkBatches.length} sources=${formatBatchSourceCounts( + sourceCounts, + )} reason=${reason}`, + { + source, + files: current.length, + chunks: chunks.length, + requests: chunkBatches.length, + sources: sourceCounts, + group: sourceWideBatchGroup, + reason, + maxFiles: SOURCE_WIDE_BATCH_MAX_FILES, + maxRequests: SOURCE_WIDE_BATCH_MAX_REQUESTS, + }, + ); + const embeddings: number[][] = []; + for (let requestIndex = 0; requestIndex < chunkBatches.length; requestIndex += 1) { + const chunkBatch = chunkBatches[requestIndex] ?? []; + embeddings.push( + ...(await this.embedChunksWithBatch(chunkBatch, firstEntry, source, { + sourceWideFiles: current.length, + sourceWideSources: sourceCounts, + sourceWideBatchGroup, + sourceWideRequestGroup: requestIndex + 1, + sourceWideRequestGroups: chunkBatches.length, + })), + ); + } + const sample = embeddings.find((embedding) => embedding.length > 0); + const vectorReady = sample ? await this.ensureVectorReady(sample.length) : false; + let offset = 0; + for (const item of current) { + const fileEmbeddings = embeddings.slice(offset, offset + item.chunks.length); + offset += item.chunks.length; + this.writeChunks( + item.entry, + item.source, + provider.model, + item.chunks, + fileEmbeddings, + vectorReady, + ); + } + prepared = []; + preparedRequestCount = 0; + }; + + for (const item of items) { + if ("kind" in item.entry && item.entry.kind === "multimodal") { + await this.indexFile(item.entry, { source: item.source }); + continue; + } + const preparedEntry = await this.prepareIndexEntry(item.entry, { source: item.source }); + if (!preparedEntry) { + continue; + } + const nextWouldExceedFiles = prepared.length >= SOURCE_WIDE_BATCH_MAX_FILES; + const nextWouldExceedRequests = + preparedRequestCount + preparedEntry.chunks.length > SOURCE_WIDE_BATCH_MAX_REQUESTS; + if (prepared.length > 0 && (nextWouldExceedFiles || nextWouldExceedRequests)) { + await flushPrepared(nextWouldExceedFiles ? "max-files" : "max-requests"); + } + prepared.push(preparedEntry); + preparedRequestCount += preparedEntry.chunks.length; + if ( + prepared.length >= SOURCE_WIDE_BATCH_MAX_FILES || + preparedRequestCount >= SOURCE_WIDE_BATCH_MAX_REQUESTS + ) { + await flushPrepared( + prepared.length >= SOURCE_WIDE_BATCH_MAX_FILES ? "max-files" : "max-requests", + ); + } + } + await flushPrepared("end"); + } + protected async indexFile( entry: MemoryIndexEntry, options: { source: MemorySource; content?: string }, @@ -757,65 +955,21 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { if ("kind" in entry && entry.kind === "multimodal") { return; } - const content = - options.content ?? - (await retryTransientMemoryRead( - () => fs.readFile(entry.absPath, "utf-8"), - `read memory markdown for indexing ${entry.absPath}`, - )); - const chunks = filterNonEmptyMemoryChunks(chunkMarkdown(content, this.settings.chunking)); - if (options.source === "sessions" && "lineMap" in entry) { - remapChunkLines(chunks, entry.lineMap); - } - this.writeChunks(entry, options.source, "fts-only", chunks, [], false); + const prepared = await this.prepareIndexEntry(entry, options); + this.writeChunks(entry, options.source, "fts-only", prepared?.chunks ?? [], [], false); return; } - let chunks: MemoryChunk[]; - let structuredInputBytes: number | undefined; - if ("kind" in entry && entry.kind === "multimodal") { - if (!this.provider) { - log.debug("Skipping multimodal indexing in FTS-only mode", { - path: entry.path, - source: options.source, - }); - this.clearIndexedFileData(entry.path, options.source); - this.upsertFileRecord(entry, options.source); - return; - } - const multimodalChunk = await buildMultimodalChunkForIndexing(entry); - if (!multimodalChunk) { - this.clearIndexedFileData(entry.path, options.source); - this.deleteFileRecord(entry.path, options.source); - return; - } - structuredInputBytes = multimodalChunk.structuredInputBytes; - chunks = [multimodalChunk.chunk]; - } else { - const content = - options.content ?? - (await retryTransientMemoryRead( - () => fs.readFile(entry.absPath, "utf-8"), - `read memory markdown for indexing ${entry.absPath}`, - )); - const baseChunks = filterNonEmptyMemoryChunks(chunkMarkdown(content, this.settings.chunking)); - chunks = this.provider - ? enforceEmbeddingMaxInputTokens(this.provider, baseChunks, EMBEDDING_BATCH_MAX_TOKENS) - : baseChunks; - if (options.source === "sessions" && "lineMap" in entry) { - remapChunkLines(chunks, entry.lineMap); - } - } - if (!this.provider) { - this.writeChunks(entry, options.source, "fts-only", chunks, [], false); + const prepared = await this.prepareIndexEntry(entry, options); + if (!prepared) { return; } let embeddings: number[][]; try { embeddings = this.batch.enabled - ? await this.embedChunksWithBatch(chunks, entry, options.source) - : await this.embedChunksInBatches(chunks); + ? await this.embedChunksWithBatch(prepared.chunks, entry, options.source) + : await this.embedChunksInBatches(prepared.chunks); } catch (err) { const message = formatErrorMessage(err); if ( @@ -827,7 +981,7 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { ) { log.warn("memory embeddings: skipping multimodal file rejected as too large", { path: entry.path, - bytes: structuredInputBytes, + bytes: prepared.structuredInputBytes, provider: this.provider.id, model: this.provider.model, error: message, @@ -840,6 +994,13 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { } const sample = embeddings.find((embedding) => embedding.length > 0); const vectorReady = sample ? await this.ensureVectorReady(sample.length) : false; - this.writeChunks(entry, options.source, this.provider.model, chunks, embeddings, vectorReady); + this.writeChunks( + entry, + options.source, + this.provider.model, + prepared.chunks, + embeddings, + vectorReady, + ); } } diff --git a/extensions/memory-core/src/memory/manager-sync-ops.ts b/extensions/memory-core/src/memory/manager-sync-ops.ts index 1b8291a393f..d6598157020 100644 --- a/extensions/memory-core/src/memory/manager-sync-ops.ts +++ b/extensions/memory-core/src/memory/manager-sync-ops.ts @@ -95,13 +95,27 @@ type MemorySyncProgressState = { report: (update: MemorySyncProgressUpdate) => void; }; -type MemoryIndexEntry = { +export type MemoryIndexEntry = { path: string; absPath: string; mtimeMs: number; size: number; hash: string; + kind?: "markdown" | "multimodal"; content?: string; + contentText?: string; + lineMap?: number[]; +}; + +export type MemoryIndexWorkItem = { + entry: MemoryIndexEntry; + source: MemorySource; + afterIndex?: () => void; +}; + +type MemorySourceSyncPlan = { + indexItems: MemoryIndexWorkItem[]; + finalize: () => Promise | void; }; const META_KEY = "memory_index_meta_v1"; @@ -111,6 +125,7 @@ const EMBEDDING_CACHE_TABLE = "embedding_cache"; const SESSION_DIRTY_DEBOUNCE_MS = 5000; const SESSION_DELTA_READ_CHUNK_BYTES = 64 * 1024; const SESSION_SYNC_YIELD_EVERY = 10; +const SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES = 128; const VECTOR_LOAD_TIMEOUT_MS = 30_000; const MEMORY_WATCH_PRESSURE_STARTUP_CHECK_DELAY_MS = 10_000; const IGNORED_MEMORY_WATCH_DIR_NAMES = new Set([ @@ -285,6 +300,96 @@ export abstract class MemoryManagerSyncOps { entry: MemoryIndexEntry, options: { source: MemorySource; content?: string }, ): Promise; + protected async indexFiles(items: MemoryIndexWorkItem[]): Promise { + for (const item of items) { + await this.indexFile(item.entry, { source: item.source }); + } + } + + private emptySourceSyncPlan(): MemorySourceSyncPlan { + return { indexItems: [], finalize: () => {} }; + } + + private shouldDeferSourceWideBatch(): boolean { + return Boolean( + this.batch.enabled && + this.provider && + this.providerRuntime?.batchEmbed && + this.providerRuntime.sourceWideBatchEmbed === true, + ); + } + + private async indexQueuedFiles( + items: MemoryIndexWorkItem[], + progress?: MemorySyncProgressState, + label?: string, + ): Promise { + if (items.length === 0) { + return; + } + if (progress && label) { + progress.report({ + completed: progress.completed, + total: progress.total, + label, + }); + } + await this.indexFiles(items); + for (const item of items) { + item.afterIndex?.(); + } + if (progress) { + progress.completed += items.length; + progress.report({ + completed: progress.completed, + total: progress.total, + }); + } + } + + private async executeSourceSyncPlans( + plans: MemorySourceSyncPlan[], + progress?: MemorySyncProgressState, + ): Promise { + const indexItems = plans.flatMap((plan) => plan.indexItems); + const sources = new Set(indexItems.map((item) => item.source)); + await this.indexQueuedFiles( + indexItems, + progress, + sources.size > 1 ? "Indexing memory sources (batch)..." : undefined, + ); + for (const plan of plans) { + await plan.finalize(); + } + } + + private async executeSourceWideSync(params: { + shouldSyncMemory: boolean; + shouldSyncSessions: boolean; + needsFullReindex: boolean; + targetSessionFiles?: string[]; + progress?: MemorySyncProgressState; + }): Promise { + const memoryPlan = params.shouldSyncMemory + ? await this.syncMemoryFiles({ + needsFullReindex: params.needsFullReindex, + progress: params.progress, + deferIndex: true, + }) + : this.emptySourceSyncPlan(); + if (params.shouldSyncSessions) { + await this.syncSessionFiles({ + needsFullReindex: params.needsFullReindex, + targetSessionFiles: params.targetSessionFiles, + progress: params.progress, + deferIndex: true, + prefixIndexItems: memoryPlan.indexItems, + }); + await memoryPlan.finalize(); + return; + } + await this.executeSourceSyncPlans([memoryPlan], params.progress); + } protected hasIndexedChunks(): boolean { const row = this.db.prepare(`SELECT 1 as found FROM chunks LIMIT 1`).get() as @@ -1504,7 +1609,8 @@ export abstract class MemoryManagerSyncOps { private async syncMemoryFiles(params: { needsFullReindex: boolean; progress?: MemorySyncProgressState; - }) { + deferIndex?: boolean; + }): Promise { const deleteFileByPathAndSource = this.db.prepare( `DELETE FROM files WHERE path = ? AND source = ?`, ); @@ -1558,8 +1664,61 @@ export abstract class MemoryManagerSyncOps { }); } - const tasks = fileEntries.map((entry) => async () => { - if (!params.needsFullReindex && existingHashes.get(entry.path) === entry.hash) { + const deleteStaleRows = async () => { + for (const stale of existingRows) { + if (activePaths.has(stale.path)) { + continue; + } + deleteFileByPathAndSource.run(stale.path, "memory"); + if (deleteVectorRowsByPathAndSource) { + try { + deleteVectorRowsByPathAndSource.run(stale.path, "memory"); + } catch {} + } + deleteChunksByPathAndSource.run(stale.path, "memory"); + if (deleteFtsRowsByPathAndSource) { + try { + deleteFtsRowsByPathAndSource.run(stale.path, "memory"); + } catch {} + } + } + }; + + if (this.batch.enabled) { + const dirtyEntries: MemoryIndexEntry[] = []; + for (const entry of fileEntries) { + if (!params.needsFullReindex && existingHashes.get(entry.path) === entry.hash) { + if (params.progress) { + params.progress.completed += 1; + params.progress.report({ + completed: params.progress.completed, + total: params.progress.total, + }); + } + continue; + } + dirtyEntries.push(entry); + } + const indexItems = dirtyEntries.map( + (entry): MemoryIndexWorkItem => ({ entry, source: "memory" }), + ); + if (params.deferIndex) { + return { indexItems, finalize: deleteStaleRows }; + } + await this.indexQueuedFiles(indexItems, params.progress); + } else { + const tasks = fileEntries.map((entry) => async () => { + if (!params.needsFullReindex && existingHashes.get(entry.path) === entry.hash) { + if (params.progress) { + params.progress.completed += 1; + params.progress.report({ + completed: params.progress.completed, + total: params.progress.total, + }); + } + return; + } + await this.indexFile(entry, { source: "memory" }); if (params.progress) { params.progress.completed += 1; params.progress.report({ @@ -1567,43 +1726,21 @@ export abstract class MemoryManagerSyncOps { total: params.progress.total, }); } - return; - } - await this.indexFile(entry, { source: "memory" }); - if (params.progress) { - params.progress.completed += 1; - params.progress.report({ - completed: params.progress.completed, - total: params.progress.total, - }); - } - }); - await runWithConcurrency(tasks, this.getIndexConcurrency()); - - for (const stale of existingRows) { - if (activePaths.has(stale.path)) { - continue; - } - deleteFileByPathAndSource.run(stale.path, "memory"); - if (deleteVectorRowsByPathAndSource) { - try { - deleteVectorRowsByPathAndSource.run(stale.path, "memory"); - } catch {} - } - deleteChunksByPathAndSource.run(stale.path, "memory"); - if (deleteFtsRowsByPathAndSource) { - try { - deleteFtsRowsByPathAndSource.run(stale.path, "memory"); - } catch {} - } + }); + await runWithConcurrency(tasks, this.getIndexConcurrency()); } + + await deleteStaleRows(); + return this.emptySourceSyncPlan(); } private async syncSessionFiles(params: { needsFullReindex: boolean; targetSessionFiles?: string[]; progress?: MemorySyncProgressState; - }) { + deferIndex?: boolean; + prefixIndexItems?: MemoryIndexWorkItem[]; + }): Promise { const deleteFileByPathAndSource = this.db.prepare( `DELETE FROM files WHERE path = ? AND source = ?`, ); @@ -1659,6 +1796,127 @@ export abstract class MemoryManagerSyncOps { } const yieldAfterSessionFile = createSessionSyncYield(files.length); + const deleteStaleRows = async () => { + if (activePaths === null) { + return; + } + + const staleRows = existingRows ?? []; + const yieldAfterStaleSessionRow = createSessionSyncYield(staleRows.length); + for (const stale of staleRows) { + try { + if (activePaths.has(stale.path)) { + continue; + } + deleteFileByPathAndSource.run(stale.path, "sessions"); + if (deleteVectorRowsByPathAndSource) { + try { + deleteVectorRowsByPathAndSource.run(stale.path, "sessions"); + } catch {} + } + deleteChunksByPathAndSource.run(stale.path, "sessions"); + if (deleteFtsRowsByPathAndSource) { + try { + deleteFtsRowsByPathAndSource.run(stale.path, "sessions"); + } catch {} + } + } finally { + await yieldAfterStaleSessionRow(); + } + } + }; + + if (params.deferIndex) { + const pendingIndexItems = [...(params.prefixIndexItems ?? [])]; + const flushPendingIndexItems = async () => { + if (pendingIndexItems.length === 0) { + return; + } + const current = pendingIndexItems.splice(0); + const sources = new Set(current.map((item) => item.source)); + await this.indexQueuedFiles( + current, + params.progress, + sources.size > 1 ? "Indexing memory sources (batch)..." : undefined, + ); + }; + + // Session entries carry flattened transcript content; flush bounded groups + // so source-wide batching cannot retain the whole dirty transcript corpus. + for (let start = 0; start < files.length; start += SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES) { + const fileBatch = files.slice(start, start + SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES); + const dirtyEntries = ( + await runWithConcurrency( + fileBatch.map((absPath) => async (): Promise => { + try { + if (!indexAll && !this.sessionsDirtyFiles.has(absPath)) { + if (params.progress) { + params.progress.completed += 1; + params.progress.report({ + completed: params.progress.completed, + total: params.progress.total, + }); + } + return null; + } + const entry = await buildSessionEntry(absPath); + if (!entry) { + if (params.progress) { + params.progress.completed += 1; + params.progress.report({ + completed: params.progress.completed, + total: params.progress.total, + }); + } + return null; + } + const existingHash = resolveMemorySourceExistingHash({ + db: this.db, + source: "sessions", + path: entry.path, + existingHashes, + }); + if (!params.needsFullReindex && existingHash === entry.hash) { + if (params.progress) { + params.progress.completed += 1; + params.progress.report({ + completed: params.progress.completed, + total: params.progress.total, + }); + } + this.resetSessionDelta(absPath, entry.size); + return null; + } + return entry; + } finally { + await yieldAfterSessionFile(); + } + }), + this.getIndexConcurrency(), + ) + ).filter((entry): entry is MemoryIndexEntry => entry !== null); + pendingIndexItems.push( + ...dirtyEntries.map( + (entry): MemoryIndexWorkItem => ({ + entry, + source: "sessions", + afterIndex: () => this.resetSessionDelta(entry.absPath, entry.size), + }), + ), + ); + if (pendingIndexItems.length >= SOURCE_WIDE_SESSION_INDEX_FLUSH_FILES) { + await flushPendingIndexItems(); + } + } + + await flushPendingIndexItems(); + await deleteStaleRows(); + return this.emptySourceSyncPlan(); + } + if ((params.prefixIndexItems?.length ?? 0) > 0) { + throw new Error("Memory session sync prefix requires deferred source-wide indexing."); + } + const tasks = files.map((absPath) => async () => { try { if (!indexAll && !this.sessionsDirtyFiles.has(absPath)) { @@ -1714,35 +1972,8 @@ export abstract class MemoryManagerSyncOps { }); await runWithConcurrency(tasks, this.getIndexConcurrency()); - if (activePaths === null) { - // Targeted syncs only refresh the requested transcripts and should not - // prune unrelated session rows without a full directory enumeration. - return; - } - - const staleRows = existingRows ?? []; - const yieldAfterStaleSessionRow = createSessionSyncYield(staleRows.length); - for (const stale of staleRows) { - try { - if (activePaths.has(stale.path)) { - continue; - } - deleteFileByPathAndSource.run(stale.path, "sessions"); - if (deleteVectorRowsByPathAndSource) { - try { - deleteVectorRowsByPathAndSource.run(stale.path, "sessions"); - } catch {} - } - deleteChunksByPathAndSource.run(stale.path, "sessions"); - if (deleteFtsRowsByPathAndSource) { - try { - deleteFtsRowsByPathAndSource.run(stale.path, "sessions"); - } catch {} - } - } finally { - await yieldAfterStaleSessionRow(); - } - } + await deleteStaleRows(); + return this.emptySourceSyncPlan(); } private createSyncProgress( @@ -1910,23 +2141,44 @@ export abstract class MemoryManagerSyncOps { ((!hasTargetSessionFiles && params?.force) || needsFullReindex || this.dirty); const shouldSyncSessions = this.shouldSyncSessions(params, needsFullReindex); - if (shouldSyncMemory) { - await this.syncMemoryFiles({ needsFullReindex, progress: progress ?? undefined }); - this.dirty = false; - } - - if (shouldSyncSessions) { - await this.syncSessionFiles({ + if (this.shouldDeferSourceWideBatch()) { + await this.executeSourceWideSync({ + shouldSyncMemory, + shouldSyncSessions, needsFullReindex, targetSessionFiles: targetSessionFiles ? Array.from(targetSessionFiles) : undefined, progress: progress ?? undefined, }); - this.sessionsDirty = false; - this.sessionsDirtyFiles.clear(); - } else if (this.sessionsDirtyFiles.size > 0) { - this.sessionsDirty = true; + if (shouldSyncMemory) { + this.dirty = false; + } + if (shouldSyncSessions) { + this.sessionsDirty = false; + this.sessionsDirtyFiles.clear(); + } else if (this.sessionsDirtyFiles.size > 0) { + this.sessionsDirty = true; + } else { + this.sessionsDirty = false; + } } else { - this.sessionsDirty = false; + if (shouldSyncMemory) { + await this.syncMemoryFiles({ needsFullReindex, progress: progress ?? undefined }); + this.dirty = false; + } + + if (shouldSyncSessions) { + await this.syncSessionFiles({ + needsFullReindex, + targetSessionFiles: targetSessionFiles ? Array.from(targetSessionFiles) : undefined, + progress: progress ?? undefined, + }); + this.sessionsDirty = false; + this.sessionsDirtyFiles.clear(); + } else if (this.sessionsDirtyFiles.size > 0) { + this.sessionsDirty = true; + } else { + this.sessionsDirty = false; + } } } catch (err) { const reason = formatErrorMessage(err); @@ -2062,6 +2314,7 @@ export abstract class MemoryManagerSyncOps { }; this.db = tempDb; + this.lastMetaSerialized = null; this.resetVectorState(); this.fts.available = false; this.fts.loadError = undefined; @@ -2087,19 +2340,39 @@ export abstract class MemoryManagerSyncOps { true, ); - if (shouldSyncMemory) { - await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress }); - this.dirty = false; - } - - if (shouldSyncSessions) { - await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress }); - this.sessionsDirty = false; - this.sessionsDirtyFiles.clear(); - } else if (this.sessionsDirtyFiles.size > 0) { - this.sessionsDirty = true; + if (this.shouldDeferSourceWideBatch()) { + await this.executeSourceWideSync({ + shouldSyncMemory, + shouldSyncSessions, + needsFullReindex: true, + progress: params.progress, + }); + if (shouldSyncMemory) { + this.dirty = false; + } + if (shouldSyncSessions) { + this.sessionsDirty = false; + this.sessionsDirtyFiles.clear(); + } else if (this.sessionsDirtyFiles.size > 0) { + this.sessionsDirty = true; + } else { + this.sessionsDirty = false; + } } else { - this.sessionsDirty = false; + if (shouldSyncMemory) { + await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress }); + this.dirty = false; + } + + if (shouldSyncSessions) { + await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress }); + this.sessionsDirty = false; + this.sessionsDirtyFiles.clear(); + } else if (this.sessionsDirtyFiles.size > 0) { + this.sessionsDirty = true; + } else { + this.sessionsDirty = false; + } } if (!shouldSyncMemory) { this.dirty = false; @@ -2170,19 +2443,39 @@ export abstract class MemoryManagerSyncOps { true, ); - if (shouldSyncMemory) { - await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress }); - this.dirty = false; - } - - if (shouldSyncSessions) { - await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress }); - this.sessionsDirty = false; - this.sessionsDirtyFiles.clear(); - } else if (this.sessionsDirtyFiles.size > 0) { - this.sessionsDirty = true; + if (this.shouldDeferSourceWideBatch()) { + await this.executeSourceWideSync({ + shouldSyncMemory, + shouldSyncSessions, + needsFullReindex: true, + progress: params.progress, + }); + if (shouldSyncMemory) { + this.dirty = false; + } + if (shouldSyncSessions) { + this.sessionsDirty = false; + this.sessionsDirtyFiles.clear(); + } else if (this.sessionsDirtyFiles.size > 0) { + this.sessionsDirty = true; + } else { + this.sessionsDirty = false; + } } else { - this.sessionsDirty = false; + if (shouldSyncMemory) { + await this.syncMemoryFiles({ needsFullReindex: true, progress: params.progress }); + this.dirty = false; + } + + if (shouldSyncSessions) { + await this.syncSessionFiles({ needsFullReindex: true, progress: params.progress }); + this.sessionsDirty = false; + this.sessionsDirtyFiles.clear(); + } else if (this.sessionsDirtyFiles.size > 0) { + this.sessionsDirty = true; + } else { + this.sessionsDirty = false; + } } if (!shouldSyncMemory) { this.dirty = false; diff --git a/extensions/openai/embedding-batch.test.ts b/extensions/openai/embedding-batch.test.ts index 5c26a8396da..e267a5cf426 100644 --- a/extensions/openai/embedding-batch.test.ts +++ b/extensions/openai/embedding-batch.test.ts @@ -1,6 +1,36 @@ // Openai tests cover embedding batch plugin behavior. -import { describe, expect, it } from "vitest"; -import { parseOpenAiBatchOutput } from "./embedding-batch.js"; +import { describe, expect, it, vi } from "vitest"; +import { parseOpenAiBatchOutput, runOpenAiEmbeddingBatches } from "./embedding-batch.js"; + +const jsonlEncoder = new TextEncoder(); + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +function jsonlBytes(value: string): number { + return jsonlEncoder.encode(value).byteLength; +} + +function fetchInputUrl(input: RequestInfo | URL): string { + if (typeof input === "string") { + return input; + } + if (input instanceof URL) { + return input.href; + } + return input.url; +} + +function parseStringBody(init: RequestInit | undefined): unknown { + if (typeof init?.body !== "string") { + throw new Error("missing JSON request body"); + } + return JSON.parse(init.body) as unknown; +} describe("OpenAI embedding batch output", () => { it("wraps malformed JSONL output", () => { @@ -8,4 +38,209 @@ describe("OpenAI embedding batch output", () => { "OpenAI embedding batch output contained malformed JSONL", ); }); + + it("splits provider uploads by serialized JSONL byte cap", async () => { + const requests: Parameters[0]["requests"] = Array.from( + { length: 3 }, + (_, index) => ({ + custom_id: String(index), + method: "POST" as const, + url: "/v1/embeddings", + body: { + model: "text-embedding-3-small", + input: `payload-${index}-${"β".repeat(8)}`, + }, + }), + ); + const uploadedJsonl: string[] = []; + const requestsByFileId = new Map>(); + const outputByFileId = new Map(); + let fileIndex = 0; + let batchIndex = 0; + const maxJsonlBytes = jsonlBytes(JSON.stringify(requests[0])); + const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = fetchInputUrl(input); + if (url.endsWith("/files") && init?.method === "POST") { + const form = init.body as FormData; + const file = form.get("file"); + if (!(file instanceof Blob)) { + throw new Error("missing batch upload file"); + } + const jsonl = await file.text(); + const fileId = `file-${fileIndex}`; + fileIndex += 1; + uploadedJsonl.push(jsonl); + requestsByFileId.set( + fileId, + jsonl.split("\n").map((line) => JSON.parse(line) as { custom_id?: string }), + ); + return jsonResponse({ id: fileId }); + } + if (url.endsWith("/batches") && init?.method === "POST") { + const body = parseStringBody(init) as { input_file_id?: string }; + const batchId = `batch-${batchIndex}`; + const outputFileId = `output-${batchIndex}`; + batchIndex += 1; + const uploadedRequests = requestsByFileId.get(body.input_file_id ?? "") ?? []; + outputByFileId.set( + outputFileId, + uploadedRequests + .map((request) => + JSON.stringify({ + custom_id: request.custom_id, + response: { + status_code: 200, + body: { data: [{ embedding: [Number(request.custom_id) + 1] }] }, + }, + }), + ) + .join("\n"), + ); + return jsonResponse({ id: batchId, status: "completed", output_file_id: outputFileId }); + } + const contentMatch = url.match(/\/files\/([^/]+)\/content$/); + if (contentMatch) { + return new Response(outputByFileId.get(contentMatch[1] ?? "") ?? "", { status: 200 }); + } + return new Response("unexpected request", { status: 500 }); + }); + + const byCustomId = await runOpenAiEmbeddingBatches({ + openAi: { + baseUrl: "https://openai-compatible.example/v1", + headers: { Authorization: "Bearer test" }, + model: "text-embedding-3-small", + fetchImpl, + }, + agentId: "main", + requests, + maxJsonlBytes, + wait: true, + concurrency: 1, + pollIntervalMs: 1000, + timeoutMs: 60_000, + }); + + expect(uploadedJsonl).toHaveLength(3); + expect(uploadedJsonl.every((jsonl) => jsonlBytes(jsonl) <= maxJsonlBytes)).toBe(true); + expect([...byCustomId.entries()]).toEqual([ + ["0", [1]], + ["1", [2]], + ["2", [3]], + ]); + }); + + it("adapts OpenAI-compatible upload groups after payload-size rejection", async () => { + const requests: Parameters[0]["requests"] = Array.from( + { length: 4 }, + (_, index) => ({ + custom_id: String(index), + method: "POST" as const, + url: "/v1/embeddings", + body: { + model: "text-embedding-3-small", + input: `payload-${index}`, + }, + }), + ); + const uploadedGroups: string[][] = []; + const requestsByFileId = new Map>(); + const outputByFileId = new Map(); + const debug = vi.fn(); + let fileIndex = 0; + let batchIndex = 0; + const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = fetchInputUrl(input); + if (url.endsWith("/files") && init?.method === "POST") { + const form = init.body as FormData; + const file = form.get("file"); + if (!(file instanceof Blob)) { + throw new Error("missing batch upload file"); + } + const uploadedRequests = (await file.text()) + .split("\n") + .map((line) => JSON.parse(line) as { custom_id?: string }); + const customIds = uploadedRequests.map((request) => request.custom_id ?? ""); + uploadedGroups.push(customIds); + if (uploadedRequests.length > 2) { + return jsonResponse( + { + error: { + message: "Request body too large. Maximum allowed: 10 MB", + type: "payload_too_large", + code: "PAYLOAD_TOO_LARGE", + }, + }, + 413, + ); + } + const fileId = `file-${fileIndex}`; + fileIndex += 1; + requestsByFileId.set(fileId, uploadedRequests); + return jsonResponse({ id: fileId }); + } + if (url.endsWith("/batches") && init?.method === "POST") { + const body = parseStringBody(init) as { input_file_id?: string }; + const batchId = `batch-${batchIndex}`; + const outputFileId = `output-${batchIndex}`; + batchIndex += 1; + const uploadedRequests = requestsByFileId.get(body.input_file_id ?? "") ?? []; + outputByFileId.set( + outputFileId, + uploadedRequests + .map((request) => + JSON.stringify({ + custom_id: request.custom_id, + response: { + status_code: 200, + body: { data: [{ embedding: [Number(request.custom_id) + 1] }] }, + }, + }), + ) + .join("\n"), + ); + return jsonResponse({ id: batchId, status: "completed", output_file_id: outputFileId }); + } + const contentMatch = url.match(/\/files\/([^/]+)\/content$/); + if (contentMatch) { + return new Response(outputByFileId.get(contentMatch[1] ?? "") ?? "", { status: 200 }); + } + return new Response("unexpected request", { status: 500 }); + }); + + const byCustomId = await runOpenAiEmbeddingBatches({ + openAi: { + baseUrl: "https://openai-compatible.example/v1", + headers: { Authorization: "Bearer test" }, + model: "text-embedding-3-small", + fetchImpl, + }, + agentId: "main", + requests, + wait: true, + concurrency: 1, + pollIntervalMs: 1000, + timeoutMs: 60_000, + debug, + }); + + expect(uploadedGroups).toEqual([ + ["0", "1", "2", "3"], + ["0", "1"], + ["2", "3"], + ]); + expect(debug).toHaveBeenCalledWith( + "memory embeddings: openai batch upload too large; splitting group", + expect.objectContaining({ + requests: 4, + parts: [2, 2], + }), + ); + expect([...byCustomId.entries()]).toEqual([ + ["0", [1]], + ["1", [2]], + ["2", [3]], + ["3", [4]], + ]); + }); }); diff --git a/extensions/openai/embedding-batch.ts b/extensions/openai/embedding-batch.ts index 5588cd97d1c..73d43b39e45 100644 --- a/extensions/openai/embedding-batch.ts +++ b/extensions/openai/embedding-batch.ts @@ -39,12 +39,22 @@ type OpenAiBatchRequest = { }; }; -type OpenAiBatchStatus = EmbeddingBatchStatus; +type OpenAiBatchStatus = EmbeddingBatchStatus & { + request_counts?: { + total?: number; + completed?: number; + failed?: number; + }; +}; type OpenAiBatchOutputLine = ProviderBatchOutputLine; export const OPENAI_BATCH_ENDPOINT = EMBEDDING_BATCH_ENDPOINT; const OPENAI_BATCH_COMPLETION_WINDOW = "24h"; const OPENAI_BATCH_MAX_REQUESTS = 50000; +// OpenAI accepts 200 MB Batch input files. Keep a safety margin so the JSONL +// splitter avoids boundary-size uploads while preserving source-wide batching. +const OPENAI_BATCH_MAX_JSONL_BYTES = 190 * 1024 * 1024; +const OPENAI_BATCH_MAX_POLL_BACKOFF_MS = 5 * 60_000; async function submitOpenAiBatch(params: { openAi: OpenAiEmbeddingClient; @@ -124,6 +134,25 @@ async function fetchOpenAiBatchResource(params: { }); } +function formatOpenAiBatchError(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function isOpenAiBatchUploadTooLargeError(error: unknown): boolean { + const message = formatOpenAiBatchError(error); + if (!/openai batch file upload failed/i.test(message)) { + return false; + } + return ( + /\b413\b/.test(message) || + /payload too large/i.test(message) || + /request body too large/i.test(message) || + /file too large/i.test(message) || + /maximum allowed/i.test(message) || + /max(?:imum)? (?:body|payload|file) (?:size )?(?:exceeded|limit)/i.test(message) + ); +} + export function parseOpenAiBatchOutput(text: string): OpenAiBatchOutputLine[] { if (!text.trim()) { return []; @@ -153,6 +182,45 @@ async function readOpenAiBatchError(params: { } } +function createOpenAiBatchPollBackoff(params: { pollIntervalMs: number; timeoutMs: number }): { + nextDelayMs: () => number; +} { + const maxDelayMs = Math.max( + params.pollIntervalMs, + Math.min(params.timeoutMs, OPENAI_BATCH_MAX_POLL_BACKOFF_MS), + ); + let delayMs = params.pollIntervalMs; + return { + nextDelayMs: () => { + const current = delayMs; + delayMs = Math.min(maxDelayMs, current * 2); + return current; + }, + }; +} + +function formatOpenAiBatchProgress(status: OpenAiBatchStatus): string { + const counts = status.request_counts; + if (!counts || typeof counts.total !== "number") { + return ""; + } + const completed = typeof counts.completed === "number" ? counts.completed : 0; + const failed = typeof counts.failed === "number" ? counts.failed : 0; + return `; progress ${completed}/${counts.total} failed=${failed}`; +} + +function formatOpenAiBatchPollError(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function isRetryableOpenAiBatchPollError(error: unknown): boolean { + const message = formatOpenAiBatchPollError(error); + return ( + /openai batch status failed: (408|409|425|429|5\d\d)\b/i.test(message) || + /\b(ECONNRESET|ECONNREFUSED|ETIMEDOUT|EAI_AGAIN)\b|fetch failed|network error/i.test(message) + ); +} + async function waitForOpenAiBatch(params: { openAi: OpenAiEmbeddingClient; batchId: string; @@ -163,14 +231,38 @@ async function waitForOpenAiBatch(params: { initial?: OpenAiBatchStatus; }): Promise { const start = Date.now(); + const pollBackoff = createOpenAiBatchPollBackoff(params); let current: OpenAiBatchStatus | undefined = params.initial; while (true) { - const status = - current ?? - (await fetchOpenAiBatchStatus({ - openAi: params.openAi, - batchId: params.batchId, - })); + let status: OpenAiBatchStatus; + try { + status = + current ?? + (await fetchOpenAiBatchStatus({ + openAi: params.openAi, + batchId: params.batchId, + })); + } catch (error) { + if (!params.wait || !isRetryableOpenAiBatchPollError(error)) { + throw error; + } + if (Date.now() - start > params.timeoutMs) { + throw new Error(`openai batch ${params.batchId} timed out after ${params.timeoutMs}ms`, { + cause: error, + }); + } + const delayMs = pollBackoff.nextDelayMs(); + params.debug?.( + `openai batch ${params.batchId} status check failed: ${formatOpenAiBatchPollError( + error, + )}; waiting ${delayMs}ms`, + ); + await new Promise((resolve) => { + setTimeout(resolve, delayMs); + }); + current = undefined; + continue; + } const state = status.status ?? "unknown"; if (state === "completed") { return resolveBatchCompletionFromStatus({ @@ -194,9 +286,14 @@ async function waitForOpenAiBatch(params: { if (Date.now() - start > params.timeoutMs) { throw new Error(`openai batch ${params.batchId} timed out after ${params.timeoutMs}ms`); } - params.debug?.(`openai batch ${params.batchId} ${state}; waiting ${params.pollIntervalMs}ms`); + const delayMs = pollBackoff.nextDelayMs(); + params.debug?.( + `openai batch ${params.batchId} ${state}${formatOpenAiBatchProgress( + status, + )}; waiting ${delayMs}ms`, + ); await new Promise((resolve) => { - setTimeout(resolve, params.pollIntervalMs); + setTimeout(resolve, delayMs); }); current = undefined; } @@ -207,13 +304,24 @@ export async function runOpenAiEmbeddingBatches( openAi: OpenAiEmbeddingClient; agentId: string; requests: OpenAiBatchRequest[]; + maxJsonlBytes?: number; } & EmbeddingBatchExecutionParams, ): Promise> { return await runEmbeddingBatchGroups({ ...buildEmbeddingBatchGroupOptions(params, { maxRequests: OPENAI_BATCH_MAX_REQUESTS, + maxJsonlBytes: params.maxJsonlBytes ?? OPENAI_BATCH_MAX_JSONL_BYTES, debugLabel: "memory embeddings: openai batch submit", }), + shouldSplitGroupOnError: isOpenAiBatchUploadTooLargeError, + onSplitGroup: ({ error, group, parts, depth }) => { + params.debug?.("memory embeddings: openai batch upload too large; splitting group", { + requests: group.length, + parts: parts.map((part) => part.length), + depth, + error: formatOpenAiBatchError(error), + }); + }, runGroup: async ({ group, groupIndex, groups, byCustomId, pollIntervalMs, timeoutMs }) => { const batchInfo = await submitOpenAiBatch({ openAi: params.openAi, diff --git a/extensions/openai/memory-embedding-adapter.ts b/extensions/openai/memory-embedding-adapter.ts index 401a740f0b4..43b41e8fccf 100644 --- a/extensions/openai/memory-embedding-adapter.ts +++ b/extensions/openai/memory-embedding-adapter.ts @@ -30,6 +30,7 @@ export const openAiMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapte provider, runtime: { id: "openai", + sourceWideBatchEmbed: true, cacheKeyData: { provider: resolvedProvider, baseUrl: client.baseUrl, diff --git a/packages/memory-host-sdk/src/host/batch-runner.test.ts b/packages/memory-host-sdk/src/host/batch-runner.test.ts index b03eeb92104..6ce3d41764c 100644 --- a/packages/memory-host-sdk/src/host/batch-runner.test.ts +++ b/packages/memory-host-sdk/src/host/batch-runner.test.ts @@ -3,6 +3,12 @@ import { describe, expect, it, vi } from "vitest"; import { MAX_SAFE_TIMEOUT_DELAY_MS } from "../../../gateway-client/src/timeouts.js"; import { buildEmbeddingBatchGroupOptions, runEmbeddingBatchGroups } from "./batch-runner.js"; +const jsonlEncoder = new TextEncoder(); + +function jsonlLineBytes(value: unknown): number { + return jsonlEncoder.encode(JSON.stringify(value)).byteLength; +} + describe("buildEmbeddingBatchGroupOptions", () => { it("clamps oversized embedding batch poll intervals to the timeout budget", () => { const options = buildEmbeddingBatchGroupOptions( @@ -61,4 +67,91 @@ describe("buildEmbeddingBatchGroupOptions", () => { expect(options.pollIntervalMs).toBe(MAX_SAFE_TIMEOUT_DELAY_MS); }); + + it("splits embedding batch groups by serialized JSONL bytes", async () => { + const requests = [ + { id: "one", body: { input: "alpha" } }, + { id: "two", body: { input: "βeta" } }, + { id: "three", body: { input: "gamma" } }, + ]; + const maxJsonlBytes = jsonlLineBytes(requests[0]) + 1 + jsonlLineBytes(requests[1]); + const groups: string[][] = []; + + await runEmbeddingBatchGroups({ + requests, + maxRequests: 100, + maxJsonlBytes, + wait: true, + pollIntervalMs: 1000, + timeoutMs: 60_000, + concurrency: 1, + debugLabel: "embedding batch submit", + runGroup: async ({ group }) => { + groups.push(group.map((request) => request.id)); + }, + }); + + expect(groups).toEqual([["one", "two"], ["three"]]); + }); + + it("splits provider-rejected batch groups when the error is splittable", async () => { + const uploadTooLarge = new Error("batch upload failed: 413 payload too large"); + const calls: string[][] = []; + const onSplitGroup = vi.fn(); + + await runEmbeddingBatchGroups({ + requests: ["one", "two", "three", "four"], + maxRequests: 100, + wait: true, + pollIntervalMs: 1000, + timeoutMs: 60_000, + concurrency: 1, + debugLabel: "embedding batch submit", + shouldSplitGroupOnError: (error) => error === uploadTooLarge, + onSplitGroup, + runGroup: async ({ group }) => { + calls.push([...group]); + if (group.length === 4) { + throw uploadTooLarge; + } + }, + }); + + expect(calls).toEqual([ + ["one", "two", "three", "four"], + ["one", "two"], + ["three", "four"], + ]); + expect(onSplitGroup).toHaveBeenCalledWith( + expect.objectContaining({ + error: uploadTooLarge, + group: ["one", "two", "three", "four"], + parts: [ + ["one", "two"], + ["three", "four"], + ], + depth: 0, + }), + ); + }); + + it("does not split a single rejected batch request", async () => { + const uploadTooLarge = new Error("batch upload failed: 413 payload too large"); + + await expect( + runEmbeddingBatchGroups({ + requests: ["one"], + maxRequests: 100, + wait: true, + pollIntervalMs: 1000, + timeoutMs: 60_000, + concurrency: 1, + debugLabel: "embedding batch submit", + shouldSplitGroupOnError: () => true, + runGroup: async () => { + throw uploadTooLarge; + }, + }), + ).rejects.toThrow(uploadTooLarge); + }); }); diff --git a/packages/memory-host-sdk/src/host/batch-runner.ts b/packages/memory-host-sdk/src/host/batch-runner.ts index 525a839925d..827edf77db4 100644 --- a/packages/memory-host-sdk/src/host/batch-runner.ts +++ b/packages/memory-host-sdk/src/host/batch-runner.ts @@ -1,6 +1,6 @@ // Memory Host SDK module implements batch runner behavior. import { resolveSafeTimeoutDelayMs } from "../../../gateway-client/src/timeouts.js"; -import { splitBatchRequests } from "./batch-utils.js"; +import { splitBatchRequestsByLimits } from "./batch-utils.js"; import { runWithConcurrency } from "./internal.js"; // Shared runner for splitting and executing remote embedding batch groups. @@ -14,6 +14,24 @@ export type EmbeddingBatchExecutionParams = { debug?: (message: string, data?: Record) => void; }; +type EmbeddingBatchGroupRunArgs = { + group: TRequest[]; + groupIndex: number; + groups: number; + byCustomId: Map; + pollIntervalMs: number; + timeoutMs: number; +}; + +type EmbeddingBatchSplitArgs = { + error: unknown; + group: TRequest[]; + parts: TRequest[][]; + groupIndex: number; + groups: number; + depth: number; +}; + /** Clamp polling to both configured poll interval and total timeout budget. */ function resolveEmbeddingBatchPollIntervalMs(params: { pollIntervalMs: number; @@ -33,41 +51,66 @@ function resolveEmbeddingBatchPollIntervalMs(params: { export async function runEmbeddingBatchGroups(params: { requests: TRequest[]; maxRequests: number; + maxJsonlBytes?: number; wait: EmbeddingBatchExecutionParams["wait"]; pollIntervalMs: EmbeddingBatchExecutionParams["pollIntervalMs"]; timeoutMs: EmbeddingBatchExecutionParams["timeoutMs"]; concurrency: EmbeddingBatchExecutionParams["concurrency"]; debugLabel: string; debug?: EmbeddingBatchExecutionParams["debug"]; - runGroup: (args: { - group: TRequest[]; - groupIndex: number; - groups: number; - byCustomId: Map; - pollIntervalMs: number; - timeoutMs: number; - }) => Promise; + shouldSplitGroupOnError?: (error: unknown, group: TRequest[]) => boolean; + onSplitGroup?: (args: EmbeddingBatchSplitArgs) => void; + runGroup: (args: EmbeddingBatchGroupRunArgs) => Promise; }): Promise> { if (params.requests.length === 0) { return new Map(); } - const groups = splitBatchRequests(params.requests, params.maxRequests); + const groups = splitBatchRequestsByLimits(params.requests, { + maxRequests: params.maxRequests, + maxJsonlBytes: params.maxJsonlBytes, + }); const byCustomId = new Map(); const pollIntervalMs = resolveEmbeddingBatchPollIntervalMs(params); + const runGroup = async (group: TRequest[], groupIndex: number, depth = 0): Promise => { + try { + await params.runGroup({ + group, + groupIndex, + groups: groups.length, + byCustomId, + pollIntervalMs, + timeoutMs: params.timeoutMs, + }); + } catch (error) { + if (group.length <= 1 || !params.shouldSplitGroupOnError?.(error, group)) { + throw error; + } + const splitAt = Math.ceil(group.length / 2); + const parts = [group.slice(0, splitAt), group.slice(splitAt)].filter( + (part) => part.length > 0, + ); + params.onSplitGroup?.({ + error, + group, + parts, + groupIndex, + groups: groups.length, + depth, + }); + for (const part of parts) { + await runGroup(part, groupIndex, depth + 1); + } + } + }; const tasks = groups.map((group, groupIndex) => async () => { - await params.runGroup({ - group, - groupIndex, - groups: groups.length, - byCustomId, - pollIntervalMs, - timeoutMs: params.timeoutMs, - }); + await runGroup(group, groupIndex); }); params.debug?.(params.debugLabel, { requests: params.requests.length, groups: groups.length, + maxRequests: params.maxRequests, + maxJsonlBytes: params.maxJsonlBytes, wait: params.wait, concurrency: params.concurrency, pollIntervalMs, @@ -81,12 +124,13 @@ export async function runEmbeddingBatchGroups(params: { /** Build normalized batch-group options for provider-specific runners. */ export function buildEmbeddingBatchGroupOptions( params: { requests: TRequest[] } & EmbeddingBatchExecutionParams, - options: { maxRequests: number; debugLabel: string }, + options: { maxRequests: number; maxJsonlBytes?: number; debugLabel: string }, ) { const pollIntervalMs = resolveEmbeddingBatchPollIntervalMs(params); return { requests: params.requests, maxRequests: options.maxRequests, + maxJsonlBytes: options.maxJsonlBytes, wait: params.wait, pollIntervalMs, timeoutMs: params.timeoutMs, diff --git a/packages/memory-host-sdk/src/host/batch-upload.ts b/packages/memory-host-sdk/src/host/batch-upload.ts index db9fbbf8666..8536984c264 100644 --- a/packages/memory-host-sdk/src/host/batch-upload.ts +++ b/packages/memory-host-sdk/src/host/batch-upload.ts @@ -30,6 +30,7 @@ export async function uploadBatchJsonlFile(params: { const filePayload = await withRemoteHttpResponse({ url: `${baseUrl}/files`, ssrfPolicy: params.client.ssrfPolicy, + fetchImpl: params.client.fetchImpl, init: { method: "POST", headers: buildBatchHeaders(params.client, { json: false }), diff --git a/packages/memory-host-sdk/src/host/batch-utils.ts b/packages/memory-host-sdk/src/host/batch-utils.ts index e205274573a..72c6f2e6981 100644 --- a/packages/memory-host-sdk/src/host/batch-utils.ts +++ b/packages/memory-host-sdk/src/host/batch-utils.ts @@ -8,6 +8,7 @@ export type BatchHttpClientConfig = { baseUrl?: string; headers?: Record; ssrfPolicy?: SsrFPolicy; + fetchImpl?: typeof fetch; }; /** Normalize batch API base URLs by removing one trailing slash. */ @@ -32,14 +33,62 @@ export function buildBatchHeaders( return headers; } +const jsonlEncoder = new TextEncoder(); + +function estimateJsonlLineBytes(request: unknown): number { + return jsonlEncoder.encode(JSON.stringify(request) ?? "").byteLength; +} + +function normalizePositiveInteger(value: number | undefined): number | undefined { + if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) { + return undefined; + } + return Math.floor(value); +} + /** Split provider requests into max-sized groups while preserving order. */ export function splitBatchRequests(requests: T[], maxRequests: number): T[][] { - if (requests.length <= maxRequests) { + const limit = normalizePositiveInteger(maxRequests) ?? 1; + if (requests.length <= limit) { return [requests]; } const groups: T[][] = []; - for (let i = 0; i < requests.length; i += maxRequests) { - groups.push(requests.slice(i, i + maxRequests)); + for (let i = 0; i < requests.length; i += limit) { + groups.push(requests.slice(i, i + limit)); + } + return groups; +} + +export function splitBatchRequestsByLimits( + requests: T[], + limits: { maxRequests: number; maxJsonlBytes?: number }, +): T[][] { + const maxRequests = normalizePositiveInteger(limits.maxRequests) ?? 1; + const maxJsonlBytes = normalizePositiveInteger(limits.maxJsonlBytes); + if (!maxJsonlBytes) { + return splitBatchRequests(requests, maxRequests); + } + + const groups: T[][] = []; + let current: T[] = []; + let currentBytes = 0; + for (const request of requests) { + const requestBytes = estimateJsonlLineBytes(request); + const separatorBytes = current.length === 0 ? 0 : 1; + const wouldExceedRequests = current.length >= maxRequests; + const wouldExceedBytes = + current.length > 0 && currentBytes + separatorBytes + requestBytes > maxJsonlBytes; + if (current.length > 0 && (wouldExceedRequests || wouldExceedBytes)) { + groups.push(current); + current = []; + currentBytes = 0; + } + + currentBytes += (current.length === 0 ? 0 : 1) + requestBytes; + current.push(request); + } + if (current.length > 0) { + groups.push(current); } return groups; } diff --git a/src/plugins/contracts/memory-embedding-provider.contract.test.ts b/src/plugins/contracts/memory-embedding-provider.contract.test.ts index 7bcfb76224f..50356382e70 100644 --- a/src/plugins/contracts/memory-embedding-provider.contract.test.ts +++ b/src/plugins/contracts/memory-embedding-provider.contract.test.ts @@ -4,7 +4,10 @@ import { registerVirtualTestPlugin, } from "openclaw/plugin-sdk/plugin-test-contracts"; import { describe, expect, it } from "vitest"; -import { getRegisteredMemoryEmbeddingProvider } from "../memory-embedding-providers.js"; +import { + getRegisteredMemoryEmbeddingProvider, + type MemoryEmbeddingBatchOptions, +} from "../memory-embedding-providers.js"; import { createPluginRecord } from "../status.test-helpers.js"; describe("memory embedding provider registration", () => { @@ -79,6 +82,58 @@ describe("memory embedding provider registration", () => { expect(provider?.ownerPluginId).toBe("memory-core"); }); + it("keeps source-wide batch embedding behind an explicit runtime opt-in", async () => { + const { config, registry } = createPluginRegistryFixture(); + + registerVirtualTestPlugin({ + registry, + config, + id: "source-wide-memory", + name: "Source Wide Memory", + contracts: { + memoryEmbeddingProviders: ["source-wide-memory"], + }, + register(api) { + api.registerMemoryEmbeddingProvider({ + id: "source-wide-memory", + create: async () => ({ + provider: { + id: "source-wide-memory", + model: "test-embedding", + embedQuery: async (text: string) => [text.length], + embedBatch: async (texts: string[]) => texts.map((text) => [text.length]), + }, + runtime: { + id: "source-wide-memory", + sourceWideBatchEmbed: true, + batchEmbed: async (batch: MemoryEmbeddingBatchOptions) => + batch.chunks.map((chunk, index) => [index, chunk.text.length]), + }, + }), + }); + }, + }); + + const adapter = getRegisteredMemoryEmbeddingProvider("source-wide-memory")?.adapter; + const result = await adapter?.create({ config, model: "test-embedding" }); + + expect(result?.runtime?.sourceWideBatchEmbed).toBe(true); + await expect( + result?.runtime?.batchEmbed?.({ + agentId: "main", + chunks: [{ text: "alpha" }, { text: "beta" }], + wait: true, + concurrency: 1, + pollIntervalMs: 1000, + timeoutMs: 60_000, + debug: () => {}, + }), + ).resolves.toEqual([ + [0, 5], + [1, 4], + ]); + }); + it("keeps companion embedding providers available during tool discovery", () => { const { config, registry } = createPluginRegistryFixture(); const record = createPluginRecord({ diff --git a/src/plugins/memory-embedding-providers.ts b/src/plugins/memory-embedding-providers.ts index 6ecdb28d1f2..5e082a06647 100644 --- a/src/plugins/memory-embedding-providers.ts +++ b/src/plugins/memory-embedding-providers.ts @@ -31,6 +31,7 @@ export type MemoryEmbeddingProviderRuntime = { cacheKeyData?: Record; inlineQueryTimeoutMs?: number; inlineBatchTimeoutMs?: number; + sourceWideBatchEmbed?: boolean; batchEmbed?: (options: MemoryEmbeddingBatchOptions) => Promise; };