Просмотр исходного кода

修复记忆合成账号饥饿并限制并行读取

Samien 1 месяц назад
Родитель
Сommit
d3f2524e67

+ 21 - 8
src/server/openVikingMemoryProvider.ts

@@ -21,6 +21,7 @@ const json = (value: unknown): string => `${JSON.stringify(value, null, 2)}\n`
 const isUuid = (value: string): boolean => /^[0-9a-f]{32}$/.test(value)
 const memoryId = (requestId: string): string => requestId.replaceAll('-', '')
 const MEMORY_TAG = 'yaoyao_memory=true'
+const MEMORY_READ_CONCURRENCY = 8
 const activeUri = (scope: MemoryScopeInput, id: string): string => `viking://~/memories/yaoyao/${scope.scope}/${id}.json`
 const revisionsUri = (scope: MemoryScopeInput, id: string, revision: number): string => `viking://resources/yaoyao/${scope.scope}/${id}/revisions/${revision}.json`
 const tombstoneUri = (scope: MemoryScopeInput, fingerprint: string): string => `viking://resources/yaoyao/${scope.scope}/tombstones/${fingerprint}.json`
@@ -84,16 +85,28 @@ export class OpenVikingMemoryProvider {
     const client = this.service.userClient(owner, scope.agentId)
     const directory = `viking://~/memories/yaoyao/${scope.scope}`
     const entries = await notFoundAsUndefined(() => client.list(directory, { recursive: true })) ?? []
-    const records: WorkspaceMemory[] = []
-    for (const entry of entries) {
+    const ids = entries.flatMap(entry => {
       const name = uriName(entry)
       const match = /([0-9a-f]{32})\.json$/.exec(name ?? '')
-      if (!match) continue
-      const record = await this.readRecord(owner, scope, match[1]!)
-      if (record && record.memory.scope === scope.scope && record.memory.agentId === scope.agentId
-        && record.memory.projectId === scope.projectId) records.push(record.memory)
-    }
-    return records
+      return match ? [match[1]!] : []
+    })
+    const records: Array<WorkspaceMemory | undefined> = new Array(ids.length)
+    let next = 0, failed = false
+    // Preserve directory order and per-record scope checks without N serial round trips.
+    const reads = await Promise.allSettled(Array.from({ length: Math.min(ids.length, MEMORY_READ_CONCURRENCY) }, async () => {
+      while (!failed) {
+        const index = next++
+        if (index >= ids.length) return
+        let record: StoredRecord | undefined
+        try { record = await this.readRecord(owner, scope, ids[index]!) }
+        catch (error) { failed = true; throw error }
+        if (record && record.memory.scope === scope.scope && record.memory.agentId === scope.agentId
+          && record.memory.projectId === scope.projectId) records[index] = record.memory
+      }
+    }))
+    const failure = reads.find(read => read.status === 'rejected')
+    if (failure?.status === 'rejected') throw failure.reason
+    return records.filter((memory): memory is WorkspaceMemory => memory !== undefined)
   }
   async list(owner: string, scopes: MemoryScopeInput[]): Promise<WorkspaceMemory[]> {
     const result: WorkspaceMemory[] = []

+ 5 - 2
src/server/workspaceMemorySynthesis.ts

@@ -39,6 +39,7 @@ export class WorkspaceMemorySynthesis {
   private timer?: ReturnType<typeof setInterval>
   private active = false
   private closed = false
+  private lastOwner?: string
   infer: (owner: string, agent: WorkspaceAgent, prompt: string) => Promise<string>
   constructor(readonly runtime: WorkspaceRuntime) {
     this.infer = async (owner, agent, prompt) => {
@@ -85,11 +86,13 @@ export class WorkspaceMemorySynthesis {
     if (this.active || this.closed) return
     this.active = true
     try {
-      for (const owner of this.runtime.store.owners()) {
+      const owners = this.runtime.store.owners()
+      const start = (owners.indexOf(this.lastOwner ?? '') + 1) % (owners.length || 1)
+      for (const owner of [...owners.slice(start), ...owners.slice(0, start)]) {
         if (this.closed || !this.runtime.userActive(owner) || owner.startsWith('_')) continue
         const jobs = this.runtime.knowledge.jobs(owner).filter(j => ['pending', 'running'].includes(j.status) && j.nextAt <= Date.now()).sort((a, b) => a.createdAt - b.createdAt)
         // One inference per tick keeps background work behind conversation work.
-        if (jobs[0]) { await this.run(owner, jobs[0]); break }
+        if (jobs[0]) { this.lastOwner = owner; await this.run(owner, jobs[0]); break }
       }
     } catch { /* Corrupt files remain untouched and are surfaced by the management API. */ }
     finally { this.active = false }

+ 51 - 0
tests/server/openVikingMemoryRead.test.ts

@@ -0,0 +1,51 @@
+// @vitest-environment node
+import { expect, it, vi } from 'vitest'
+import { OpenVikingMemoryProvider } from '../../src/server/openVikingMemoryProvider'
+import type { OpenVikingService } from '../../src/server/openVikingService'
+import type { WorkspaceKnowledge } from '../../src/server/workspaceKnowledge'
+
+function fixture() {
+  const records = Array.from({ length: 20 }, (_, index) => ({
+    schemaVersion: 1, revisions: [], memory: {
+      id: index.toString(16).padStart(32, '0'), scope: 'agent', agentId: 'agent', content: `fact-${index}`,
+    },
+  }))
+  const list = vi.fn(async () => records.map(record => `${record.memory.id}.json`))
+  const read = vi.fn(async (uri: string) => JSON.stringify(records.find(record => uri.endsWith(`${record.memory.id}.json`))))
+  const client = { list, read }
+  const service = { enabled: true, ensureUser: vi.fn(async () => undefined), userClient: vi.fn(() => client) }
+  const knowledge = { agent: vi.fn(() => ({ id: 'agent' })) }
+  return { provider: new OpenVikingMemoryProvider(service as unknown as OpenVikingService, knowledge as unknown as WorkspaceKnowledge),
+    records, read, list }
+}
+
+it('overlaps at most eight reads while preserving listing order and scope isolation', async () => {
+  const f = fixture()
+  let active = 0, peak = 0
+  const release: Array<() => void> = []
+  f.records[3]!.memory.agentId = 'another-agent'
+  f.read.mockImplementation(async uri => {
+    active++; peak = Math.max(peak, active)
+    await new Promise<void>(resolve => release.push(resolve))
+    active--
+    return JSON.stringify(f.records.find(record => uri.endsWith(`${record.memory.id}.json`)))
+  })
+  const pending = f.provider.list('owner', [{ scope: 'agent', agentId: 'agent' }])
+  await vi.waitFor(() => expect(release).toHaveLength(8))
+  // Release batches in reverse order, forcing responses to arrive out of listing order.
+  while (f.read.mock.calls.length < 20 || active) {
+    release.splice(0).reverse().forEach(resolve => resolve())
+    await new Promise(resolve => setTimeout(resolve, 0))
+  }
+  expect((await pending).map(memory => memory.content)).toEqual(f.records.filter(record => record.memory.agentId === 'agent').map(record => record.memory.content))
+  expect(peak).toBe(8)
+  expect(f.read).toHaveBeenCalledTimes(20)
+})
+
+it('keeps upstream failures visible and continues to reject corrupt records', async () => {
+  const f = fixture()
+  f.read.mockRejectedValueOnce(new Error('upstream unavailable'))
+  await expect(f.provider.list('owner', [{ scope: 'agent', agentId: 'agent' }])).rejects.toThrow('upstream unavailable')
+  f.read.mockResolvedValueOnce('invalid JSON')
+  await expect(f.provider.list('owner', [{ scope: 'agent', agentId: 'agent' }])).rejects.toMatchObject({ code: 'openviking_memory_corrupt' })
+})

+ 21 - 0
tests/server/workspaceMemorySynthesis.test.ts

@@ -30,6 +30,27 @@ function turn(content = '以后我喜欢简洁的中文回复') {
   synthesis.enqueue(owner, run)
   return { run, source: store.require<WorkspaceMessage>(owner, 'message', run.messageId), reply }
 }
+it('rotates accounts even when the first account always has eligible memory work', async () => {
+  vi.spyOn(store, 'owners').mockReturnValue(['first', 'second', 'inactive', '_internal'])
+  vi.spyOn(runtime, 'userActive').mockImplementation(account => account !== 'inactive')
+  vi.spyOn(runtime.knowledge, 'jobs').mockImplementation(account => [{
+    id: `${account}-job`, status: 'pending', nextAt: 0, createdAt: account === 'first' ? 100 : 0,
+  }] as any)
+  const run = vi.spyOn(synthesis, 'run').mockResolvedValue()
+  for (let index = 0; index < 6; index++) await synthesis.tick()
+  expect(run.mock.calls.map(([account]) => account)).toEqual(['first', 'second', 'first', 'second', 'first', 'second'])
+})
+it('does not reserve a fairness turn for accounts without due jobs and handles a removed account', async () => {
+  const owners = vi.spyOn(store, 'owners').mockReturnValue(['first', 'second'])
+  vi.spyOn(runtime.knowledge, 'jobs').mockImplementation(account => [{
+    id: `${account}-job`, status: 'pending', nextAt: account === 'first' ? Date.now() + 60_000 : 0, createdAt: 0,
+  }] as any)
+  const run = vi.spyOn(synthesis, 'run').mockResolvedValue()
+  await synthesis.tick()
+  owners.mockReturnValue(['third'])
+  await synthesis.tick()
+  expect(run.mock.calls.map(([account]) => account)).toEqual(['second', 'third'])
+})
 it('runs extraction on server Hermes even for a legacy computer Bot',async()=>{
   const request=vi.fn(async()=>({status:200,body:Buffer.from(JSON.stringify({text:'{"memories":[]}'}))}))
   const target={session:{request}}