stack/apps/gateway/src/agent/tools/memory-tools.ts

import { Type } from '@sinclair/typebox';
import type { ToolDefinition } from '@mariozechner/pi-coding-agent';
import type { Memory } from '@mosaic/memory';
import type { EmbeddingProvider } from '@mosaic/memory';

/**
 * Create memory tools bound to the session's authenticated userId.
 *
 * SECURITY: userId is resolved from the authenticated session at tool-creation
 * time and is never accepted as a user-supplied or LLM-supplied parameter.
 * This prevents cross-user data access via parameter injection.
 */
export function createMemoryTools(
  memory: Memory,
  embeddingProvider: EmbeddingProvider | null,
  /** Authenticated user ID from the session. All memory operations are scoped to this user. */
  sessionUserId: string | undefined,
): ToolDefinition[] {
  /** Return an error result when no session user is bound. */
  function noUserError() {
    return {
      content: [
        {
          type: 'text' as const,
          text: 'Memory tools unavailable — no authenticated user bound to this session',
        },
      ],
      details: undefined,
    };
  }

  const searchMemory: ToolDefinition = {
    name: 'memory_search',
    label: 'Search Memory',
    description:
      'Search across stored insights and knowledge using natural language. Returns semantically similar results.',
    parameters: Type.Object({
      query: Type.String({ description: 'Natural language search query' }),
      limit: Type.Optional(Type.Number({ description: 'Max results (default 5)' })),
    }),
    async execute(_toolCallId, params) {
      if (!sessionUserId) return noUserError();

      const { query, limit } = params as {
        query: string;
        limit?: number;
      };

      if (!embeddingProvider) {
        return {
          content: [
            {
              type: 'text' as const,
              text: 'Semantic search unavailable — no embedding provider configured',
            },
          ],
          details: undefined,
        };
      }

      const embedding = await embeddingProvider.embed(query);
      const results = await memory.insights.searchByEmbedding(sessionUserId, embedding, limit ?? 5);
      return {
        content: [{ type: 'text' as const, text: JSON.stringify(results, null, 2) }],
        details: undefined,
      };
    },
  };

  const getPreferences: ToolDefinition = {
    name: 'memory_get_preferences',
    label: 'Get User Preferences',
    description: 'Retrieve stored preferences for the current session user.',
    parameters: Type.Object({
      category: Type.Optional(
        Type.String({
          description: 'Filter by category: communication, coding, workflow, appearance, general',
        }),
      ),
    }),
    async execute(_toolCallId, params) {
      if (!sessionUserId) return noUserError();

      const { category } = params as { category?: string };
      type Cat = 'communication' | 'coding' | 'workflow' | 'appearance' | 'general';
      const prefs = category
        ? await memory.preferences.findByUserAndCategory(sessionUserId, category as Cat)
        : await memory.preferences.findByUser(sessionUserId);
      return {
        content: [{ type: 'text' as const, text: JSON.stringify(prefs, null, 2) }],
        details: undefined,
      };
    },
  };

  const savePreference: ToolDefinition = {
    name: 'memory_save_preference',
    label: 'Save User Preference',
    description:
      'Store a learned user preference (e.g., "prefers tables over paragraphs", "timezone: America/Chicago").',
    parameters: Type.Object({
      key: Type.String({ description: 'Preference key' }),
      value: Type.String({ description: 'Preference value (JSON string)' }),
      category: Type.Optional(
        Type.String({
          description: 'Category: communication, coding, workflow, appearance, general',
        }),
      ),
    }),
    async execute(_toolCallId, params) {
      if (!sessionUserId) return noUserError();

      const { key, value, category } = params as {
        key: string;
        value: string;
        category?: string;
      };
      type Cat = 'communication' | 'coding' | 'workflow' | 'appearance' | 'general';
      let parsedValue: unknown;
      try {
        parsedValue = JSON.parse(value);
      } catch {
        parsedValue = value;
      }
      const pref = await memory.preferences.upsert({
        userId: sessionUserId,
        key,
        value: parsedValue,
        category: (category as Cat) ?? 'general',
        source: 'agent',
      });
      return {
        content: [{ type: 'text' as const, text: JSON.stringify(pref, null, 2) }],
        details: undefined,
      };
    },
  };

  const saveInsight: ToolDefinition = {
    name: 'memory_save_insight',
    label: 'Save Insight',
    description:
      'Store a learned insight, decision, or knowledge extracted from the current interaction.',
    parameters: Type.Object({
      content: Type.String({ description: 'The insight or knowledge to store' }),
      category: Type.Optional(
        Type.String({
          description: 'Category: decision, learning, preference, fact, pattern, general',
        }),
      ),
    }),
    async execute(_toolCallId, params) {
      if (!sessionUserId) return noUserError();

      const { content, category } = params as {
        content: string;
        category?: string;
      };
      type Cat = 'decision' | 'learning' | 'preference' | 'fact' | 'pattern' | 'general';

      let embedding: number[] | null = null;
      if (embeddingProvider) {
        embedding = await embeddingProvider.embed(content);
      }

      const insight = await memory.insights.create({
        userId: sessionUserId,
        content,
        embedding,
        source: 'agent',
        category: (category as Cat) ?? 'learning',
      });
      return {
        content: [{ type: 'text' as const, text: JSON.stringify(insight, null, 2) }],
        details: undefined,
      };
    },
  };

  return [searchMemory, getPreferences, savePreference, saveInsight];
}