From 31a59738089f0784428833fc5a0192c6c7c43261 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Mon, 13 Jul 2026 08:18:14 -0500 Subject: [PATCH] feat(memory): bind operator plugin to agent sessions --- apps/gateway/src/agent/agent.service.ts | 23 +++++++++-- .../src/agent/tools/memory-tools.test.ts | 41 +++++++++++++++++++ apps/gateway/src/agent/tools/memory-tools.ts | 32 +++++++++++++-- apps/gateway/src/memory/memory.module.ts | 22 +++++++++- 4 files changed, 111 insertions(+), 7 deletions(-) create mode 100644 apps/gateway/src/agent/tools/memory-tools.test.ts diff --git a/apps/gateway/src/agent/agent.service.ts b/apps/gateway/src/agent/agent.service.ts index f5b2347..4ac42d5 100644 --- a/apps/gateway/src/agent/agent.service.ts +++ b/apps/gateway/src/agent/agent.service.ts @@ -15,9 +15,10 @@ import { type ToolDefinition, } from '@mariozechner/pi-coding-agent'; import type { Brain } from '@mosaicstack/brain'; -import type { Memory } from '@mosaicstack/memory'; +import type { Memory, OperatorMemoryPlugin } from '@mosaicstack/memory'; import { BRAIN } from '../brain/brain.tokens.js'; import { MEMORY } from '../memory/memory.tokens.js'; +import { OPERATOR_MEMORY_PLUGIN } from '../memory/memory.module.js'; import { EmbeddingService } from '../memory/embedding.service.js'; import { CoordService } from '../coord/coord.service.js'; import { ProviderService } from './provider.service.js'; @@ -135,6 +136,9 @@ export class AgentService implements OnModuleDestroy { @Inject(PreferencesService) private readonly preferencesService: PreferencesService | null, @Inject(SessionGCService) private readonly gc: SessionGCService, + @Optional() + @Inject(OPERATOR_MEMORY_PLUGIN) + private readonly operatorMemory: OperatorMemoryPlugin | null = null, ) {} /** @@ -146,6 +150,7 @@ export class AgentService implements OnModuleDestroy { private buildToolsForSandbox( sandboxDir: string, sessionUserId: string | undefined, + sessionScope?: { tenantId: string; ownerId: string; sessionId: string }, ): ToolDefinition[] { return [ ...createBrainTools(this.brain), @@ -154,6 +159,9 @@ export class AgentService implements OnModuleDestroy { this.memory, this.embeddingService.available ? this.embeddingService : null, sessionUserId, + this.operatorMemory && sessionScope + ? { plugin: this.operatorMemory, scope: sessionScope } + : undefined, ), ...createFileTools(sandboxDir), ...createGitTools(sandboxDir), @@ -228,6 +236,7 @@ export class AgentService implements OnModuleDestroy { isAdmin: options.isAdmin, agentConfigId: options.agentConfigId, userId: options.userId, + tenantId: options.tenantId, conversationHistory: options.conversationHistory, }; this.logger.log( @@ -267,7 +276,15 @@ export class AgentService implements OnModuleDestroy { } // Build per-session tools scoped to the sandbox directory and authenticated user - const sandboxTools = this.buildToolsForSandbox(sandboxDir, mergedOptions?.userId); + const sessionUserId = mergedOptions?.userId; + const sessionTenantId = this.tenantIdFor(sessionUserId, mergedOptions?.tenantId); + const sandboxTools = this.buildToolsForSandbox( + sandboxDir, + sessionUserId, + sessionUserId && sessionTenantId + ? { tenantId: sessionTenantId, ownerId: sessionUserId, sessionId } + : undefined, + ); // Combine static tools with dynamically discovered MCP client tools and skill tools const mcpTools = this.mcpClientService.getToolDefinitions(); @@ -362,7 +379,7 @@ export class AgentService implements OnModuleDestroy { sandboxDir, allowedTools, userId: mergedOptions?.userId, - tenantId: this.tenantIdFor(mergedOptions?.userId, mergedOptions?.tenantId), + tenantId: sessionTenantId, agentConfigId: mergedOptions?.agentConfigId, agentName: resolvedAgentName, metrics: { diff --git a/apps/gateway/src/agent/tools/memory-tools.test.ts b/apps/gateway/src/agent/tools/memory-tools.test.ts new file mode 100644 index 0000000..d90e7bf --- /dev/null +++ b/apps/gateway/src/agent/tools/memory-tools.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it, vi } from 'vitest'; +import { createMemoryTools } from './memory-tools.js'; + +describe('createMemoryTools operator retrieval binding', () => { + const memory = { + insights: { searchByEmbedding: vi.fn(), create: vi.fn() }, + preferences: { findByUserAndCategory: vi.fn(), findByUser: vi.fn(), upsert: vi.fn() }, + }; + const scope = { tenantId: 'tenant-a', ownerId: 'owner-a', sessionId: 'session-a' }; + + it('uses the configured plugin with the server-derived scope for retrieval and capture', async () => { + const plugin = { + search: vi.fn(async () => []), + capture: vi.fn(async () => ({ id: 'insight-1' })), + }; + const tools = createMemoryTools(memory as never, null, 'owner-a', { + plugin: plugin as never, + scope, + }); + + await tools + .find((tool) => tool.name === 'memory_search')! + .execute('call-1', { query: 'plans' }, undefined, undefined, {} as never); + await tools + .find((tool) => tool.name === 'memory_save_insight')! + .execute( + 'call-2', + { content: 'secret', category: 'decision' }, + undefined, + undefined, + {} as never, + ); + + expect(plugin.search).toHaveBeenCalledWith(scope, 'plans', 5); + expect(plugin.capture).toHaveBeenCalledWith(scope, { + content: 'secret', + source: 'agent', + category: 'decision', + }); + }); +}); diff --git a/apps/gateway/src/agent/tools/memory-tools.ts b/apps/gateway/src/agent/tools/memory-tools.ts index ab1809a..ec9b744 100644 --- a/apps/gateway/src/agent/tools/memory-tools.ts +++ b/apps/gateway/src/agent/tools/memory-tools.ts @@ -1,7 +1,11 @@ import { Type } from '@sinclair/typebox'; import type { ToolDefinition } from '@mariozechner/pi-coding-agent'; -import type { Memory } from '@mosaicstack/memory'; -import type { EmbeddingProvider } from '@mosaicstack/memory'; +import type { + EmbeddingProvider, + Memory, + OperatorMemoryPlugin, + OperatorMemoryScope, +} from '@mosaicstack/memory'; /** * Create memory tools bound to the session's authenticated userId. @@ -13,8 +17,10 @@ import type { EmbeddingProvider } from '@mosaicstack/memory'; export function createMemoryTools( memory: Memory, embeddingProvider: EmbeddingProvider | null, - /** Authenticated user ID from the session. All memory operations are scoped to this user. */ + /** Authenticated user ID from the session. All preference operations are scoped to this user. */ sessionUserId: string | undefined, + /** Optional configured retrieval plugin, bound to a server-derived session scope. */ + operatorMemory?: { plugin: OperatorMemoryPlugin; scope: OperatorMemoryScope }, ): ToolDefinition[] { /** Return an error result when no session user is bound. */ function noUserError() { @@ -46,6 +52,14 @@ export function createMemoryTools( limit?: number; }; + if (operatorMemory) { + const results = await operatorMemory.plugin.search(operatorMemory.scope, query, limit ?? 5); + return { + content: [{ type: 'text' as const, text: JSON.stringify(results, null, 2) }], + details: undefined, + }; + } + if (!embeddingProvider) { return { content: [ @@ -158,6 +172,18 @@ export function createMemoryTools( }; type Cat = 'decision' | 'learning' | 'preference' | 'fact' | 'pattern' | 'general'; + if (operatorMemory) { + const insight = await operatorMemory.plugin.capture(operatorMemory.scope, { + content, + source: 'agent', + category: category ?? 'learning', + }); + return { + content: [{ type: 'text' as const, text: JSON.stringify(insight, null, 2) }], + details: undefined, + }; + } + let embedding: number[] | null = null; if (embeddingProvider) { embedding = await embeddingProvider.embed(content); diff --git a/apps/gateway/src/memory/memory.module.ts b/apps/gateway/src/memory/memory.module.ts index 779ad40..56e0278 100644 --- a/apps/gateway/src/memory/memory.module.ts +++ b/apps/gateway/src/memory/memory.module.ts @@ -3,8 +3,10 @@ import { createMemory, type Memory, createMemoryAdapter, + createOperatorMemoryPlugin, type MemoryAdapter, type MemoryConfig, + type OperatorMemoryPlugin, } from '@mosaicstack/memory'; import type { Db } from '@mosaicstack/db'; import type { StorageAdapter } from '@mosaicstack/storage'; @@ -14,6 +16,9 @@ import { DB, STORAGE_ADAPTER } from '../database/database.module.js'; import { MEMORY } from './memory.tokens.js'; import { MemoryController } from './memory.controller.js'; import { EmbeddingService } from './embedding.service.js'; +import { redactSensitiveContent } from '@mosaicstack/log'; + +export const OPERATOR_MEMORY_PLUGIN = 'OPERATOR_MEMORY_PLUGIN'; export const MEMORY_ADAPTER = 'MEMORY_ADAPTER'; @@ -38,9 +43,24 @@ function buildMemoryConfig(config: MosaicConfig, storageAdapter: StorageAdapter) createMemoryAdapter(buildMemoryConfig(config, storageAdapter)), inject: [MOSAIC_CONFIG, STORAGE_ADAPTER], }, + { + provide: OPERATOR_MEMORY_PLUGIN, + useFactory: (adapter: MemoryAdapter): OperatorMemoryPlugin | null => { + const instanceId = process.env['MOSAIC_OPERATOR_MEMORY_INSTANCE_ID']?.trim(); + const namespace = process.env['MOSAIC_OPERATOR_MEMORY_NAMESPACE']?.trim(); + if (!instanceId || !namespace) return null; + return createOperatorMemoryPlugin({ + adapter, + instanceId, + namespace, + redact: (content) => redactSensitiveContent(content).content, + }); + }, + inject: [MEMORY_ADAPTER], + }, EmbeddingService, ], controllers: [MemoryController], - exports: [MEMORY, MEMORY_ADAPTER, EmbeddingService], + exports: [MEMORY, MEMORY_ADAPTER, OPERATOR_MEMORY_PLUGIN, EmbeddingService], }) export class MemoryModule {}