From 4fefaa8a78900e8f77a38c056506abb921f6790c Mon Sep 17 00:00:00 2001 From: Jack Chen Date: Thu, 23 Apr 2026 18:45:48 +0800 Subject: [PATCH] feat: expose custom MemoryStore via TeamConfig.sharedMemoryStore (#157) * feat: expose custom MemoryStore via TeamConfig.sharedMemoryStore (#156) The MemoryStore interface was already public, and SharedMemory / TeamInfo already used the interface internally. This adds the final user-config wire so integrators can attach custom backends (Redis, Postgres, Engram, etc.) without hacking SharedMemory private fields. Priority: sharedMemoryStore > sharedMemory: true > no memory. Fully backward-compatible: existing sharedMemory: true users see no change. SDK-only: the CLI cannot pass runtime objects through JSON config. Closes #156. * fix: validate sharedMemoryStore shape and reject from CLI JSON Addresses Codex P2 review on #157. Plain objects from untrusted JSON could previously reach `new SharedMemory(plainObject)` and crash later on the first `.set`/`.list` call with a cryptic TypeError. Defense-in-depth: - SharedMemory constructor performs a runtime shape check and throws a clear TypeError if the provided store does not implement get/set/list/delete/clear. - CLI `asTeamConfig` explicitly rejects `sharedMemoryStore` in JSON with a message pointing to the SDK path, since this field is documented SDK-only. Adds 4 tests covering malformed stores (plain object, partial interface, null, Team constructor path). * fix: route falsy-but-present sharedMemoryStore through shape check Addresses Codex finding 2 on #157. The truthy gate silently skipped falsy values (null, 0, ''), letting config bugs downgrade to the default in-memory store or no memory instead of failing fast. Switched to `!== undefined` so any present value reaches SharedMemory's runtime shape check and throws a clear TypeError. Adds 2 tests: null as a bogus store throws; omitting the field still honors `sharedMemory: true`. --- CLAUDE.md | 2 +- README.md | 28 ++++++ README_zh.md | 28 ++++++ docs/cli.md | 2 + src/cli/oma.ts | 11 +++ src/memory/shared.ts | 40 +++++++- src/team/team.ts | 13 ++- src/types.ts | 9 ++ tests/shared-memory.test.ts | 178 ++++++++++++++++++++++++++++++++++++ 9 files changed, 306 insertions(+), 5 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index a586a50..4956488 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -46,7 +46,7 @@ This is the framework's key feature. When `runTeam()` is called: | Task | `task/queue.ts`, `task/task.ts` | Dependency-aware queue, auto-unblock on completion, cascade failure to dependents | | Tool | `tool/framework.ts`, `tool/executor.ts`, `tool/built-in/` | `defineTool()` with Zod schemas, ToolRegistry, parallel batch execution with concurrency semaphore | | LLM | `llm/adapter.ts`, `llm/anthropic.ts`, `llm/openai.ts` | `LLMAdapter` interface (`chat` + `stream`), factory `createAdapter()` | -| Memory | `memory/shared.ts`, `memory/store.ts` | Namespaced key-value store (`agentName/key`), markdown summary injection into prompts | +| Memory | `memory/shared.ts`, `memory/store.ts` | Namespaced key-value store (`agentName/key`), markdown summary injection into prompts. Custom backends via `TeamConfig.sharedMemoryStore` (any `MemoryStore` impl); `sharedMemory: true` uses the default in-process store | | Types | `types.ts` | All interfaces in one file to avoid circular deps | | Exports | `index.ts` | Public API surface | diff --git a/README.md b/README.md index a936bbf..acce83f 100644 --- a/README.md +++ b/README.md @@ -333,6 +333,34 @@ Notes: See [`integrations/mcp-github`](examples/integrations/mcp-github.ts) for a full runnable setup. +## Shared Memory + +Teams can share a namespaced key-value store so later agents see earlier agents' findings. Enable it with a boolean for the default in-process store: + +```typescript +const team = orchestrator.createTeam('research-team', { + name: 'research-team', + agents: [researcher, writer], + sharedMemory: true, +}) +``` + +For durable or cross-process backends (Redis, Postgres, Engram, etc.), implement the `MemoryStore` interface and pass it via `sharedMemoryStore`. Keys are still namespaced as `/` before reaching the store: + +```typescript +import type { MemoryStore } from '@jackchen_me/open-multi-agent' + +class RedisStore implements MemoryStore { /* get/set/list/delete/clear */ } + +const team = orchestrator.createTeam('durable-team', { + name: 'durable-team', + agents: [researcher, writer], + sharedMemoryStore: new RedisStore(), +}) +``` + +When both are provided, `sharedMemoryStore` wins. SDK-only: the CLI cannot pass runtime objects. + ## Context Management Long-running agents can hit input token ceilings fast. Set `contextStrategy` on `AgentConfig` to control how the conversation shrinks as it grows: diff --git a/README_zh.md b/README_zh.md index 8384e0f..0098f13 100644 --- a/README_zh.md +++ b/README_zh.md @@ -331,6 +331,34 @@ await disconnect() 完整例子见 [`integrations/mcp-github`](examples/integrations/mcp-github.ts)。 +## 共享内存 + +团队可以共用一个命名空间化的 key-value 存储，让后续 agent 看到前面 agent 的发现。用布尔值启用默认的进程内存储： + +```typescript +const team = orchestrator.createTeam('research-team', { + name: 'research-team', + agents: [researcher, writer], + sharedMemory: true, +}) +``` + +需要持久化或跨进程的后端（Redis、Postgres、Engram 等）？实现 `MemoryStore` 接口并通过 `sharedMemoryStore` 注入，键仍会在到达 store 前按 `/` 做命名空间封装： + +```typescript +import type { MemoryStore } from '@jackchen_me/open-multi-agent' + +class RedisStore implements MemoryStore { /* get/set/list/delete/clear */ } + +const team = orchestrator.createTeam('durable-team', { + name: 'durable-team', + agents: [researcher, writer], + sharedMemoryStore: new RedisStore(), +}) +``` + +两者都提供时，`sharedMemoryStore` 优先。此字段仅 SDK 可用，CLI 无法序列化运行时对象。 + ## 上下文管理 长时间运行的 agent 很容易撞上输入 token 上限。在 `AgentConfig` 里设 `contextStrategy`，控制对话变长时怎么收缩： diff --git a/docs/cli.md b/docs/cli.md index 77a670d..4175de5 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -121,6 +121,8 @@ Validation rules enforced by the CLI: Any other fields are passed through to the library as in TypeScript. +**SDK-only fields**: `sharedMemoryStore` (custom `MemoryStore` instance) cannot be set from JSON since it is a runtime object. Use `sharedMemory: true` for the default in-memory store, or wire a custom store in TypeScript via `orchestrator.createTeam()`. + ### Tasks file Used with **`oma task --file`**. diff --git a/src/cli/oma.ts b/src/cli/oma.ts index 04d7fdf..1548024 100644 --- a/src/cli/oma.ts +++ b/src/cli/oma.ts @@ -141,6 +141,17 @@ function asTeamConfig(v: unknown, label: string): TeamConfig { throw new OmaValidationError(`agent.model required for "${String(a['name'])}"`) } } + // `sharedMemoryStore` is a runtime MemoryStore instance and cannot survive + // JSON round-tripping. Reject it here with a clear pointer to the SDK path, + // otherwise the plain object would reach `new SharedMemory(...)` and crash on + // the first read/write. + if ('sharedMemoryStore' in v) { + throw new OmaValidationError( + `${label}.sharedMemoryStore: SDK-only; cannot be set from JSON config. ` + + 'Use `sharedMemory: true` for the default in-memory store, or wire a ' + + 'custom MemoryStore in TypeScript via `orchestrator.createTeam()`.', + ) + } return v as unknown as TeamConfig } diff --git a/src/memory/shared.ts b/src/memory/shared.ts index 00649c2..f7cc85b 100644 --- a/src/memory/shared.ts +++ b/src/memory/shared.ts @@ -10,6 +10,25 @@ import type { MemoryEntry, MemoryStore } from '../types.js' import { InMemoryStore } from './store.js' +// --------------------------------------------------------------------------- +// Runtime shape check +// --------------------------------------------------------------------------- + +const STORE_METHODS = ['get', 'set', 'list', 'delete', 'clear'] as const + +/** + * Returns true when `v` structurally implements {@link MemoryStore}. + * + * Used to defend against malformed `sharedMemoryStore` values reaching + * {@link SharedMemory} (e.g. a plain object deserialized from JSON that + * cannot actually satisfy the interface at runtime). + */ +function isMemoryStore(v: unknown): v is MemoryStore { + if (v === null || typeof v !== 'object') return false + const obj = v as Record + return STORE_METHODS.every((m) => typeof obj[m] === 'function') +} + // --------------------------------------------------------------------------- // SharedMemory // --------------------------------------------------------------------------- @@ -34,10 +53,25 @@ import { InMemoryStore } from './store.js' * ``` */ export class SharedMemory { - private readonly store: InMemoryStore + private readonly store: MemoryStore - constructor() { - this.store = new InMemoryStore() + /** + * @param store - Optional custom {@link MemoryStore} backing this shared memory. + * Defaults to an in-process {@link InMemoryStore}. Custom stores + * receive namespaced keys (`/`) opaque to them. + * + * @throws {TypeError} when `store` is provided but does not structurally + * implement {@link MemoryStore} (fails fast on malformed + * values, e.g. plain objects from untrusted JSON config). + */ + constructor(store?: MemoryStore) { + if (store !== undefined && !isMemoryStore(store)) { + throw new TypeError( + 'SharedMemory: `store` must implement the MemoryStore interface ' + + `(methods: ${STORE_METHODS.join(', ')}).`, + ) + } + this.store = store ?? new InMemoryStore() } // --------------------------------------------------------------------------- diff --git a/src/team/team.ts b/src/team/team.ts index c88148b..864fc78 100644 --- a/src/team/team.ts +++ b/src/team/team.ts @@ -103,7 +103,18 @@ export class Team { this.agentMap = new Map(config.agents.map((a) => [a.name, a])) this.bus = new MessageBus() this.queue = new TaskQueue() - this.memory = config.sharedMemory ? new SharedMemory() : undefined + // Resolve shared memory: + // - `sharedMemoryStore` takes precedence when present (enables memory regardless of boolean). + // - `sharedMemory: true` with no custom store → default in-memory store. + // - otherwise → no shared memory. + // Use `!== undefined` rather than a truthy check so that malformed falsy + // values (null, 0, '') still reach SharedMemory's shape validation and + // fail fast, instead of silently falling back and hiding the config bug. + this.memory = config.sharedMemoryStore !== undefined + ? new SharedMemory(config.sharedMemoryStore) + : config.sharedMemory + ? new SharedMemory() + : undefined this.events = new EventBus() // Bridge queue events onto the team's event bus. diff --git a/src/types.ts b/src/types.ts index fa60ca0..c4e4572 100644 --- a/src/types.ts +++ b/src/types.ts @@ -435,6 +435,15 @@ export interface TeamConfig { readonly name: string readonly agents: readonly AgentConfig[] readonly sharedMemory?: boolean + /** + * Custom {@link MemoryStore} backing the team's shared memory (e.g. Redis, + * Postgres, or a remote service). When provided, shared memory is enabled + * regardless of `sharedMemory`. When both are set, `sharedMemoryStore` wins. + * When omitted and `sharedMemory` is `true`, the default in-memory store is used. + * + * SDK-only: the CLI (`oma`) cannot pass runtime objects through its JSON config. + */ + readonly sharedMemoryStore?: MemoryStore readonly maxConcurrency?: number } diff --git a/tests/shared-memory.test.ts b/tests/shared-memory.test.ts index 87f28d8..89ace1f 100644 --- a/tests/shared-memory.test.ts +++ b/tests/shared-memory.test.ts @@ -1,5 +1,7 @@ import { describe, it, expect } from 'vitest' import { SharedMemory } from '../src/memory/shared.js' +import { Team } from '../src/team/team.js' +import type { MemoryEntry, MemoryStore } from '../src/types.js' describe('SharedMemory', () => { // ------------------------------------------------------------------------- @@ -132,4 +134,180 @@ describe('SharedMemory', () => { const all = await mem.listAll() expect(all).toHaveLength(2) }) + + // ------------------------------------------------------------------------- + // Custom MemoryStore injection (issue #156) + // ------------------------------------------------------------------------- + + describe('custom MemoryStore injection', () => { + /** Recording store that forwards to an internal map and tracks every call. */ + class RecordingStore implements MemoryStore { + readonly data = new Map() + readonly setCalls: Array<{ key: string; value: string }> = [] + + async get(key: string): Promise { + return this.data.get(key) ?? null + } + async set( + key: string, + value: string, + metadata?: Record, + ): Promise { + this.setCalls.push({ key, value }) + this.data.set(key, { key, value, metadata, createdAt: new Date() }) + } + async list(): Promise { + return Array.from(this.data.values()) + } + async delete(key: string): Promise { + this.data.delete(key) + } + async clear(): Promise { + this.data.clear() + } + } + + it('routes writes through an injected MemoryStore', async () => { + const store = new RecordingStore() + const mem = new SharedMemory(store) + await mem.write('alice', 'plan', 'v1') + + expect(store.setCalls).toEqual([{ key: 'alice/plan', value: 'v1' }]) + }) + + it('preserves `/` namespace prefix on the underlying store', async () => { + const store = new RecordingStore() + const mem = new SharedMemory(store) + await mem.write('bob', 'notes', 'hello') + + const entry = await store.get('bob/notes') + expect(entry?.value).toBe('hello') + }) + + it('getSummary reads from the injected store', async () => { + const store = new RecordingStore() + const mem = new SharedMemory(store) + await mem.write('alice', 'k', 'val') + + const summary = await mem.getSummary() + expect(summary).toContain('### alice') + expect(summary).toContain('k: val') + }) + + it('getStore returns the injected store', () => { + const store = new RecordingStore() + const mem = new SharedMemory(store) + expect(mem.getStore()).toBe(store) + }) + + it('Team wires `sharedMemoryStore` into its SharedMemory', async () => { + const store = new RecordingStore() + const team = new Team({ + name: 'injection-team', + agents: [{ name: 'alice', model: 'claude-sonnet-4-6' }], + sharedMemoryStore: store, + }) + + const sharedMem = team.getSharedMemoryInstance() + expect(sharedMem).toBeDefined() + await sharedMem!.write('alice', 'fact', 'committed') + + expect(store.setCalls).toEqual([{ key: 'alice/fact', value: 'committed' }]) + }) + + it('Team: `sharedMemoryStore` takes precedence over `sharedMemory: false`', () => { + const store = new RecordingStore() + const team = new Team({ + name: 'override-team', + agents: [{ name: 'alice', model: 'claude-sonnet-4-6' }], + sharedMemory: false, + sharedMemoryStore: store, + }) + + // Custom store wins: memory is enabled even though the boolean is false. + expect(team.getSharedMemoryInstance()).toBeDefined() + expect(team.getSharedMemory()).toBe(store) + }) + + it('Team: neither flag → no shared memory (backward compat)', () => { + const team = new Team({ + name: 'no-memory-team', + agents: [{ name: 'alice', model: 'claude-sonnet-4-6' }], + }) + expect(team.getSharedMemoryInstance()).toBeUndefined() + }) + + it('Team: `sharedMemory: true` only → default InMemoryStore (backward compat)', () => { + const team = new Team({ + name: 'default-memory-team', + agents: [{ name: 'alice', model: 'claude-sonnet-4-6' }], + sharedMemory: true, + }) + expect(team.getSharedMemoryInstance()).toBeDefined() + expect(team.getSharedMemory()).toBeDefined() + }) + + // ----------------------------------------------------------------------- + // Shape validation — defends against malformed `sharedMemoryStore` + // (e.g. plain objects from untrusted JSON) reaching SharedMemory. + // ----------------------------------------------------------------------- + + it('SharedMemory throws when store is a plain object missing methods', () => { + const plain = { foo: 'bar' } as unknown as MemoryStore + expect(() => new SharedMemory(plain)).toThrow(TypeError) + expect(() => new SharedMemory(plain)).toThrow(/MemoryStore interface/) + }) + + it('SharedMemory throws when store is missing a single method', () => { + const partial = { + get: async () => null, + set: async () => undefined, + list: async () => [], + delete: async () => undefined, + // `clear` missing + } as unknown as MemoryStore + expect(() => new SharedMemory(partial)).toThrow(TypeError) + }) + + it('SharedMemory throws when store is null (cast)', () => { + expect(() => new SharedMemory(null as unknown as MemoryStore)).toThrow(TypeError) + }) + + it('Team throws early on malformed `sharedMemoryStore`', () => { + const bogus = { not: 'a store' } as unknown as MemoryStore + expect( + () => + new Team({ + name: 'bad-team', + agents: [{ name: 'alice', model: 'claude-sonnet-4-6' }], + sharedMemoryStore: bogus, + }), + ).toThrow(TypeError) + }) + + it('Team throws on falsy-but-present sharedMemoryStore (null)', () => { + // `null` is falsy but present; a truthy gate would silently drop it. + // The `!== undefined` gate routes it through SharedMemory's shape check + // so config bugs fail fast instead of being silently downgraded. + expect( + () => + new Team({ + name: 'null-store-team', + agents: [{ name: 'alice', model: 'claude-sonnet-4-6' }], + sharedMemoryStore: null as unknown as MemoryStore, + }), + ).toThrow(TypeError) + }) + + it('Team: omitting sharedMemoryStore entirely still honors sharedMemory: true', () => { + // Sanity check that the `!== undefined` gate does not accidentally + // enable memory when the field is absent. + const team = new Team({ + name: 'absent-store-team', + agents: [{ name: 'alice', model: 'claude-sonnet-4-6' }], + sharedMemory: true, + }) + expect(team.getSharedMemoryInstance()).toBeDefined() + }) + }) })