feat(memory): bind operator plugin to agent sessions (#739)

apps/gateway/src/agent/__tests__/agent-service-ownership.test.ts

@@ -12,18 +12,24 @@ type AgentServiceInternals = {
   creating: Map<string, Promise<AgentSession>>;
 };
 
-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,
+    { getToolDefinitions: vi.fn(() => []) } as never,
-    {} 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();

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: {

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',
+    });
+  });
+});

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 {
-import type { EmbeddingProvider } from '@mosaicstack/memory';
+  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);

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 {}