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

Merged
jason.woltje merged 2 commits from fix/tess-m4-operator-memory-consumer into main 2026-07-13 13:44:21 +00:00
4 changed files with 111 additions and 7 deletions
Showing only changes of commit 31a5973808 - Show all commits

View File

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

View File

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

View File

@@ -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);

View File

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