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/prisma/migrations/20260228000000_ms22_agent_memory/migration.sql

@@ -0,0 +1,24 @@
+-- CreateTable
+CREATE TABLE "agent_memories" (
+    "id" UUID NOT NULL,
+    "workspace_id" UUID NOT NULL,
+    "agent_id" TEXT NOT NULL,
+    "key" TEXT NOT NULL,
+    "value" JSONB NOT NULL,
+    "created_at" TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMPTZ NOT NULL,
+
+    CONSTRAINT "agent_memories_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX "agent_memories_workspace_id_agent_id_key_key" ON "agent_memories"("workspace_id", "agent_id", "key");
+
+-- CreateIndex
+CREATE INDEX "agent_memories_workspace_id_idx" ON "agent_memories"("workspace_id");
+
+-- CreateIndex
+CREATE INDEX "agent_memories_agent_id_idx" ON "agent_memories"("agent_id");
+
+-- AddForeignKey
+ALTER TABLE "agent_memories" ADD CONSTRAINT "agent_memories_workspace_id_fkey" FOREIGN KEY ("workspace_id") REFERENCES "workspaces"("id") ON DELETE CASCADE ON UPDATE CASCADE;

apps/api/prisma/schema.prisma

@@ -299,6 +299,7 @@ model Workspace {
   agentSessions                AgentSession[]
   agentTasks                   AgentTask[]
   findings                     Finding[]
+  agentMemories                AgentMemory[]
   userLayouts                  UserLayout[]
   knowledgeEntries             KnowledgeEntry[]
   knowledgeTags                KnowledgeTag[]
@@ -764,6 +765,23 @@ model AgentSession {
   @@map("agent_sessions")
 }
 
+model AgentMemory {
+  id          String   @id @default(uuid()) @db.Uuid
+  workspaceId String   @map("workspace_id") @db.Uuid
+  agentId     String   @map("agent_id")
+  key         String
+  value       Json
+  createdAt   DateTime @default(now()) @map("created_at") @db.Timestamptz
+  updatedAt   DateTime @updatedAt @map("updated_at") @db.Timestamptz
+
+  workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
+
+  @@unique([workspaceId, agentId, key])
+  @@index([workspaceId])
+  @@index([agentId])
+  @@map("agent_memories")
+}
+
 model WidgetDefinition {
   id String @id @default(uuid()) @db.Uuid
 

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,

docs/scratchpads/ms22-agent-memory.md

@@ -0,0 +1,64 @@
+# MS22 Agent Memory Module
+
+## Objective
+
+Add per-agent key/value store: AgentMemory model + NestJS module with CRUD endpoints.
+
+## Issues
+
+- MS22-DB-002: Add AgentMemory schema model
+- MS22-API-002: Add agent-memory NestJS module
+
+## Plan
+
+1. AgentMemory model → schema.prisma (after AgentSession, line 736)
+2. Add `agentMemories AgentMemory[]` relation to Workspace model
+3. Create apps/api/src/agent-memory/ with service, controller, DTOs, specs
+4. Register in app.module.ts
+5. Migrate: `prisma migrate dev --name ms22_agent_memory`
+6. lint + build
+7. Commit
+
+## Endpoints
+
+- PUT /api/agents/:agentId/memory/:key (upsert)
+- GET /api/agents/:agentId/memory (list all)
+- GET /api/agents/:agentId/memory/:key (get one)
+- DELETE /api/agents/:agentId/memory/:key (remove)
+
+## Auth
+
+- @UseGuards(AuthGuard, WorkspaceGuard, PermissionGuard)
+- @Workspace() decorator for workspaceId
+- Permission.WORKSPACE_MEMBER for write ops
+- Permission.WORKSPACE_ANY for read ops
+
+## Schema
+
+```prisma
+model AgentMemory {
+  id          String   @id @default(uuid()) @db.Uuid
+  workspaceId String   @map("workspace_id") @db.Uuid
+  agentId     String   @map("agent_id")
+  key         String
+  value       Json
+  createdAt   DateTime @default(now()) @map("created_at") @db.Timestamptz
+  updatedAt   DateTime @updatedAt @map("updated_at") @db.Timestamptz
+
+  workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
+
+  @@unique([workspaceId, agentId, key])
+  @@index([workspaceId])
+  @@index([agentId])
+  @@map("agent_memories")
+}
+```
+
+## Progress
+
+- [ ] Schema
+- [ ] Module files
+- [ ] app.module.ts
+- [ ] Migration
+- [ ] lint/build
+- [ ] Commit