diff --git a/apps/gateway/src/agent/__tests__/agent-service-ownership.test.ts b/apps/gateway/src/agent/__tests__/agent-service-ownership.test.ts index b015155..83e0368 100644 --- a/apps/gateway/src/agent/__tests__/agent-service-ownership.test.ts +++ b/apps/gateway/src/agent/__tests__/agent-service-ownership.test.ts @@ -12,18 +12,24 @@ type AgentServiceInternals = { creating: Map>; }; -function makeService(): AgentService { +function makeService(operatorMemory: unknown = null): AgentService { return new AgentService( - {} as never, + { + getDefaultModel: vi.fn(() => null), + getRegistry: vi.fn(() => ({})), + findModel: vi.fn(), + listAvailableModels: vi.fn(() => []), + } as never, {} as never, {} as never, { available: false } as never, {} as never, - {} as never, - {} as never, + { getToolDefinitions: vi.fn(() => []) } as never, + { loadForSession: vi.fn(async () => ({ metaTools: [], promptAdditions: [] })) } as never, null, null, { collect: vi.fn().mockResolvedValue(undefined) } as never, + operatorMemory as never, ); } @@ -120,6 +126,37 @@ describe('AgentService owner/tenant scope enforcement', () => { expect(internals(service).sessions.has(CONVERSATION_ID)).toBe(false); }); + it('derives the operator-memory scope on the createSession production path', async () => { + const plugin = { capture: vi.fn(), search: vi.fn() }; + const service = makeService(plugin); + const buildTools = vi.spyOn(service as never, 'buildToolsForSandbox').mockReturnValue([]); + + // Session construction reaches the real scope derivation before the intentionally incomplete + // Pi test double rejects later in createAgentSession. + await service.createSession(CONVERSATION_ID, OWNER_SCOPE).catch(() => undefined); + + expect(buildTools).toHaveBeenCalledWith(expect.any(String), OWNER_SCOPE.userId, { + tenantId: OWNER_SCOPE.tenantId, + ownerId: OWNER_SCOPE.userId, + sessionId: CONVERSATION_ID, + }); + }); + + it('denies a foreign actor before it can obtain another session operator-memory scope', async () => { + const plugin = { capture: vi.fn(), search: vi.fn() }; + const service = makeService(plugin); + internals(service).sessions.set(CONVERSATION_ID, makeSession()); + const buildTools = vi.spyOn(service as never, 'buildToolsForSandbox'); + + await expect(service.createSession(CONVERSATION_ID, FOREIGN_SCOPE)).rejects.toBeInstanceOf( + ForbiddenException, + ); + + expect(buildTools).not.toHaveBeenCalled(); + expect(plugin.capture).not.toHaveBeenCalled(); + expect(plugin.search).not.toHaveBeenCalled(); + }); + it('checks owner/tenant scope before returning an in-flight session creation', async () => { const service = makeService(); const session = makeSession(); 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 {}