feat(api): add agent memory module (MS22-DB-002, MS22-API-002) (#586)

Co-authored-by: Jason Woltje <jason@diversecanvas.com>
Co-committed-by: Jason Woltje <jason@diversecanvas.com>

apps/api/src/agent-memory/agent-memory.controller.spec.ts

@@ -0,0 +1,102 @@
+import { Test, TestingModule } from "@nestjs/testing";
+import { AgentMemoryController } from "./agent-memory.controller";
+import { AgentMemoryService } from "./agent-memory.service";
+import { AuthGuard } from "../auth/guards/auth.guard";
+import { WorkspaceGuard, PermissionGuard } from "../common/guards";
+import { describe, it, expect, beforeEach, vi } from "vitest";
+
+describe("AgentMemoryController", () => {
+  let controller: AgentMemoryController;
+
+  const mockAgentMemoryService = {
+    upsert: vi.fn(),
+    findAll: vi.fn(),
+    findOne: vi.fn(),
+    remove: vi.fn(),
+  };
+
+  const mockGuard = { canActivate: vi.fn(() => true) };
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      controllers: [AgentMemoryController],
+      providers: [
+        {
+          provide: AgentMemoryService,
+          useValue: mockAgentMemoryService,
+        },
+      ],
+    })
+      .overrideGuard(AuthGuard)
+      .useValue(mockGuard)
+      .overrideGuard(WorkspaceGuard)
+      .useValue(mockGuard)
+      .overrideGuard(PermissionGuard)
+      .useValue(mockGuard)
+      .compile();
+
+    controller = module.get<AgentMemoryController>(AgentMemoryController);
+
+    vi.clearAllMocks();
+  });
+
+  const workspaceId = "workspace-1";
+  const agentId = "agent-1";
+  const key = "context";
+
+  describe("upsert", () => {
+    it("should upsert a memory entry", async () => {
+      const dto = { value: { foo: "bar" } };
+      const mockEntry = { id: "mem-1", workspaceId, agentId, key, value: dto.value };
+
+      mockAgentMemoryService.upsert.mockResolvedValue(mockEntry);
+
+      const result = await controller.upsert(agentId, key, dto, workspaceId);
+
+      expect(mockAgentMemoryService.upsert).toHaveBeenCalledWith(workspaceId, agentId, key, dto);
+      expect(result).toEqual(mockEntry);
+    });
+  });
+
+  describe("findAll", () => {
+    it("should list all memory entries for an agent", async () => {
+      const mockEntries = [
+        { id: "mem-1", key: "a", value: 1 },
+        { id: "mem-2", key: "b", value: 2 },
+      ];
+
+      mockAgentMemoryService.findAll.mockResolvedValue(mockEntries);
+
+      const result = await controller.findAll(agentId, workspaceId);
+
+      expect(mockAgentMemoryService.findAll).toHaveBeenCalledWith(workspaceId, agentId);
+      expect(result).toEqual(mockEntries);
+    });
+  });
+
+  describe("findOne", () => {
+    it("should get a single memory entry", async () => {
+      const mockEntry = { id: "mem-1", key, value: "v" };
+
+      mockAgentMemoryService.findOne.mockResolvedValue(mockEntry);
+
+      const result = await controller.findOne(agentId, key, workspaceId);
+
+      expect(mockAgentMemoryService.findOne).toHaveBeenCalledWith(workspaceId, agentId, key);
+      expect(result).toEqual(mockEntry);
+    });
+  });
+
+  describe("remove", () => {
+    it("should delete a memory entry", async () => {
+      const mockResponse = { message: "Memory entry deleted successfully" };
+
+      mockAgentMemoryService.remove.mockResolvedValue(mockResponse);
+
+      const result = await controller.remove(agentId, key, workspaceId);
+
+      expect(mockAgentMemoryService.remove).toHaveBeenCalledWith(workspaceId, agentId, key);
+      expect(result).toEqual(mockResponse);
+    });
+  });
+});

apps/api/src/agent-memory/agent-memory.controller.ts

@@ -0,0 +1,89 @@
+import {
+  Controller,
+  Get,
+  Put,
+  Delete,
+  Body,
+  Param,
+  UseGuards,
+  HttpCode,
+  HttpStatus,
+} from "@nestjs/common";
+import { AgentMemoryService } from "./agent-memory.service";
+import { UpsertAgentMemoryDto } from "./dto";
+import { AuthGuard } from "../auth/guards/auth.guard";
+import { WorkspaceGuard, PermissionGuard } from "../common/guards";
+import { Workspace, Permission, RequirePermission } from "../common/decorators";
+
+/**
+ * Controller for per-agent key/value memory endpoints.
+ * All endpoints require authentication and workspace context.
+ *
+ * Guards are applied in order:
+ * 1. AuthGuard - Verifies user authentication
+ * 2. WorkspaceGuard - Validates workspace access
+ * 3. PermissionGuard - Checks role-based permissions
+ */
+@Controller("agents/:agentId/memory")
+@UseGuards(AuthGuard, WorkspaceGuard, PermissionGuard)
+export class AgentMemoryController {
+  constructor(private readonly agentMemoryService: AgentMemoryService) {}
+
+  /**
+   * PUT /api/agents/:agentId/memory/:key
+   * Upsert a memory entry for an agent
+   * Requires: MEMBER role or higher
+   */
+  @Put(":key")
+  @RequirePermission(Permission.WORKSPACE_MEMBER)
+  async upsert(
+    @Param("agentId") agentId: string,
+    @Param("key") key: string,
+    @Body() dto: UpsertAgentMemoryDto,
+    @Workspace() workspaceId: string
+  ) {
+    return this.agentMemoryService.upsert(workspaceId, agentId, key, dto);
+  }
+
+  /**
+   * GET /api/agents/:agentId/memory
+   * List all memory entries for an agent
+   * Requires: Any workspace member (including GUEST)
+   */
+  @Get()
+  @RequirePermission(Permission.WORKSPACE_ANY)
+  async findAll(@Param("agentId") agentId: string, @Workspace() workspaceId: string) {
+    return this.agentMemoryService.findAll(workspaceId, agentId);
+  }
+
+  /**
+   * GET /api/agents/:agentId/memory/:key
+   * Get a single memory entry by key
+   * Requires: Any workspace member (including GUEST)
+   */
+  @Get(":key")
+  @RequirePermission(Permission.WORKSPACE_ANY)
+  async findOne(
+    @Param("agentId") agentId: string,
+    @Param("key") key: string,
+    @Workspace() workspaceId: string
+  ) {
+    return this.agentMemoryService.findOne(workspaceId, agentId, key);
+  }
+
+  /**
+   * DELETE /api/agents/:agentId/memory/:key
+   * Remove a memory entry
+   * Requires: MEMBER role or higher
+   */
+  @Delete(":key")
+  @HttpCode(HttpStatus.OK)
+  @RequirePermission(Permission.WORKSPACE_MEMBER)
+  async remove(
+    @Param("agentId") agentId: string,
+    @Param("key") key: string,
+    @Workspace() workspaceId: string
+  ) {
+    return this.agentMemoryService.remove(workspaceId, agentId, key);
+  }
+}

apps/api/src/agent-memory/agent-memory.module.ts

@@ -0,0 +1,13 @@
+import { Module } from "@nestjs/common";
+import { AgentMemoryController } from "./agent-memory.controller";
+import { AgentMemoryService } from "./agent-memory.service";
+import { PrismaModule } from "../prisma/prisma.module";
+import { AuthModule } from "../auth/auth.module";
+
+@Module({
+  imports: [PrismaModule, AuthModule],
+  controllers: [AgentMemoryController],
+  providers: [AgentMemoryService],
+  exports: [AgentMemoryService],
+})
+export class AgentMemoryModule {}

apps/api/src/agent-memory/agent-memory.service.spec.ts

@@ -0,0 +1,126 @@
+import { Test, TestingModule } from "@nestjs/testing";
+import { AgentMemoryService } from "./agent-memory.service";
+import { PrismaService } from "../prisma/prisma.service";
+import { NotFoundException } from "@nestjs/common";
+import { describe, it, expect, beforeEach, vi } from "vitest";
+
+describe("AgentMemoryService", () => {
+  let service: AgentMemoryService;
+
+  const mockPrismaService = {
+    agentMemory: {
+      upsert: vi.fn(),
+      findMany: vi.fn(),
+      findUnique: vi.fn(),
+      delete: vi.fn(),
+    },
+  };
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [
+        AgentMemoryService,
+        {
+          provide: PrismaService,
+          useValue: mockPrismaService,
+        },
+      ],
+    }).compile();
+
+    service = module.get<AgentMemoryService>(AgentMemoryService);
+
+    vi.clearAllMocks();
+  });
+
+  const workspaceId = "workspace-1";
+  const agentId = "agent-1";
+  const key = "session-context";
+
+  describe("upsert", () => {
+    it("should upsert a memory entry", async () => {
+      const dto = { value: { data: "some context" } };
+      const mockEntry = {
+        id: "mem-1",
+        workspaceId,
+        agentId,
+        key,
+        value: dto.value,
+        createdAt: new Date(),
+        updatedAt: new Date(),
+      };
+
+      mockPrismaService.agentMemory.upsert.mockResolvedValue(mockEntry);
+
+      const result = await service.upsert(workspaceId, agentId, key, dto);
+
+      expect(mockPrismaService.agentMemory.upsert).toHaveBeenCalledWith({
+        where: { workspaceId_agentId_key: { workspaceId, agentId, key } },
+        create: { workspaceId, agentId, key, value: dto.value },
+        update: { value: dto.value },
+      });
+      expect(result).toEqual(mockEntry);
+    });
+  });
+
+  describe("findAll", () => {
+    it("should return all memory entries for an agent", async () => {
+      const mockEntries = [
+        { id: "mem-1", key: "a", value: 1 },
+        { id: "mem-2", key: "b", value: 2 },
+      ];
+
+      mockPrismaService.agentMemory.findMany.mockResolvedValue(mockEntries);
+
+      const result = await service.findAll(workspaceId, agentId);
+
+      expect(mockPrismaService.agentMemory.findMany).toHaveBeenCalledWith({
+        where: { workspaceId, agentId },
+        orderBy: { key: "asc" },
+      });
+      expect(result).toEqual(mockEntries);
+    });
+  });
+
+  describe("findOne", () => {
+    it("should return a memory entry by key", async () => {
+      const mockEntry = { id: "mem-1", workspaceId, agentId, key, value: "ctx" };
+
+      mockPrismaService.agentMemory.findUnique.mockResolvedValue(mockEntry);
+
+      const result = await service.findOne(workspaceId, agentId, key);
+
+      expect(mockPrismaService.agentMemory.findUnique).toHaveBeenCalledWith({
+        where: { workspaceId_agentId_key: { workspaceId, agentId, key } },
+      });
+      expect(result).toEqual(mockEntry);
+    });
+
+    it("should throw NotFoundException when key not found", async () => {
+      mockPrismaService.agentMemory.findUnique.mockResolvedValue(null);
+
+      await expect(service.findOne(workspaceId, agentId, key)).rejects.toThrow(NotFoundException);
+    });
+  });
+
+  describe("remove", () => {
+    it("should delete a memory entry", async () => {
+      const mockEntry = { id: "mem-1", workspaceId, agentId, key, value: "x" };
+
+      mockPrismaService.agentMemory.findUnique.mockResolvedValue(mockEntry);
+      mockPrismaService.agentMemory.delete.mockResolvedValue(mockEntry);
+
+      const result = await service.remove(workspaceId, agentId, key);
+
+      expect(mockPrismaService.agentMemory.delete).toHaveBeenCalledWith({
+        where: { workspaceId_agentId_key: { workspaceId, agentId, key } },
+      });
+      expect(result).toEqual({ message: "Memory entry deleted successfully" });
+    });
+
+    it("should throw NotFoundException when key not found", async () => {
+      mockPrismaService.agentMemory.findUnique.mockResolvedValue(null);
+
+      await expect(service.remove(workspaceId, agentId, key)).rejects.toThrow(NotFoundException);
+    });
+  });
+});

apps/api/src/agent-memory/agent-memory.service.ts

@@ -0,0 +1,79 @@
+import { Injectable, NotFoundException } from "@nestjs/common";
+import { PrismaService } from "../prisma/prisma.service";
+import { Prisma } from "@prisma/client";
+import type { UpsertAgentMemoryDto } from "./dto";
+
+@Injectable()
+export class AgentMemoryService {
+  constructor(private readonly prisma: PrismaService) {}
+
+  /**
+   * Upsert a memory entry for an agent.
+   */
+  async upsert(workspaceId: string, agentId: string, key: string, dto: UpsertAgentMemoryDto) {
+    return this.prisma.agentMemory.upsert({
+      where: {
+        workspaceId_agentId_key: { workspaceId, agentId, key },
+      },
+      create: {
+        workspaceId,
+        agentId,
+        key,
+        value: dto.value as Prisma.InputJsonValue,
+      },
+      update: {
+        value: dto.value as Prisma.InputJsonValue,
+      },
+    });
+  }
+
+  /**
+   * List all memory entries for an agent in a workspace.
+   */
+  async findAll(workspaceId: string, agentId: string) {
+    return this.prisma.agentMemory.findMany({
+      where: { workspaceId, agentId },
+      orderBy: { key: "asc" },
+    });
+  }
+
+  /**
+   * Get a single memory entry by key.
+   */
+  async findOne(workspaceId: string, agentId: string, key: string) {
+    const entry = await this.prisma.agentMemory.findUnique({
+      where: {
+        workspaceId_agentId_key: { workspaceId, agentId, key },
+      },
+    });
+
+    if (!entry) {
+      throw new NotFoundException(`Memory key "${key}" not found for agent "${agentId}"`);
+    }
+
+    return entry;
+  }
+
+  /**
+   * Delete a memory entry by key.
+   */
+  async remove(workspaceId: string, agentId: string, key: string) {
+    const entry = await this.prisma.agentMemory.findUnique({
+      where: {
+        workspaceId_agentId_key: { workspaceId, agentId, key },
+      },
+    });
+
+    if (!entry) {
+      throw new NotFoundException(`Memory key "${key}" not found for agent "${agentId}"`);
+    }
+
+    await this.prisma.agentMemory.delete({
+      where: {
+        workspaceId_agentId_key: { workspaceId, agentId, key },
+      },
+    });
+
+    return { message: "Memory entry deleted successfully" };
+  }
+}

apps/api/src/agent-memory/dto/index.ts

@@ -0,0 +1 @@
+export * from "./upsert-agent-memory.dto";

apps/api/src/agent-memory/dto/upsert-agent-memory.dto.ts

@@ -0,0 +1,10 @@
+import { IsNotEmpty } from "class-validator";
+
+/**
+ * DTO for upserting an agent memory entry.
+ * The value accepts any JSON-serializable data.
+ */
+export class UpsertAgentMemoryDto {
+  @IsNotEmpty({ message: "value must not be empty" })
+  value!: unknown;
+}

apps/api/src/app.module.ts

@@ -28,6 +28,7 @@ import { BrainModule } from "./brain/brain.module";
 import { CronModule } from "./cron/cron.module";
 import { AgentTasksModule } from "./agent-tasks/agent-tasks.module";
 import { FindingsModule } from "./findings/findings.module";
+import { AgentMemoryModule } from "./agent-memory/agent-memory.module";
 import { ValkeyModule } from "./valkey/valkey.module";
 import { BullMqModule } from "./bullmq/bullmq.module";
 import { StitcherModule } from "./stitcher/stitcher.module";
@@ -102,6 +103,7 @@ import { RlsContextInterceptor } from "./common/interceptors/rls-context.interce
     CronModule,
     AgentTasksModule,
     FindingsModule,
+    AgentMemoryModule,
     RunnerJobsModule,
     JobEventsModule,
     JobStepsModule,