feat: tool registration — brain tools for agent sessions (P2-004) (#76)

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

apps/gateway/package.json

@@ -30,6 +30,7 @@
     "@opentelemetry/sdk-metrics": "^2.6.0",
     "@opentelemetry/sdk-node": "^0.213.0",
     "@opentelemetry/semantic-conventions": "^1.40.0",
+    "@sinclair/typebox": "^0.34.48",
     "better-auth": "^1.5.5",
     "fastify": "^5.0.0",
     "reflect-metadata": "^0.2.0",

apps/gateway/src/agent/agent.service.ts

@@ -1,11 +1,15 @@
-import { Injectable, Logger, type OnModuleDestroy } from '@nestjs/common';
+import { Inject, Injectable, Logger, type OnModuleDestroy } from '@nestjs/common';
 import {
   createAgentSession,
   SessionManager,
   type AgentSession as PiAgentSession,
   type AgentSessionEvent,
+  type ToolDefinition,
 } from '@mariozechner/pi-coding-agent';
+import type { Brain } from '@mosaic/brain';
+import { BRAIN } from '../brain/brain.tokens.js';
 import { ProviderService } from './provider.service.js';
+import { createBrainTools } from './tools/brain-tools.js';
 
 export interface AgentSessionOptions {
   provider?: string;
@@ -27,7 +31,15 @@ export class AgentService implements OnModuleDestroy {
   private readonly sessions = new Map<string, AgentSession>();
   private readonly creating = new Map<string, Promise<AgentSession>>();
 
-  constructor(private readonly providerService: ProviderService) {}
+  private readonly customTools: ToolDefinition[];
+
+  constructor(
+    private readonly providerService: ProviderService,
+    @Inject(BRAIN) private readonly brain: Brain,
+  ) {
+    this.customTools = createBrainTools(brain);
+    this.logger.log(`Registered ${this.customTools.length} custom tools`);
+  }
 
   async createSession(sessionId: string, options?: AgentSessionOptions): Promise<AgentSession> {
     const existing = this.sessions.get(sessionId);
@@ -62,6 +74,7 @@ export class AgentService implements OnModuleDestroy {
         modelRegistry: this.providerService.getRegistry(),
         model: model ?? undefined,
         tools: [],
+        customTools: this.customTools,
       });
       piSession = result.session;
     } catch (err) {

apps/gateway/src/agent/tools/brain-tools.ts

@@ -0,0 +1,184 @@
+import { Type } from '@sinclair/typebox';
+import type { ToolDefinition } from '@mariozechner/pi-coding-agent';
+import type { Brain } from '@mosaic/brain';
+
+export function createBrainTools(brain: Brain): ToolDefinition[] {
+  const listProjects: ToolDefinition = {
+    name: 'brain_list_projects',
+    label: 'List Projects',
+    description: 'List all projects in the brain.',
+    parameters: Type.Object({}),
+    async execute() {
+      const projects = await brain.projects.findAll();
+      return {
+        content: [{ type: 'text' as const, text: JSON.stringify(projects, null, 2) }],
+        details: undefined,
+      };
+    },
+  };
+
+  const getProject: ToolDefinition = {
+    name: 'brain_get_project',
+    label: 'Get Project',
+    description: 'Get a project by ID.',
+    parameters: Type.Object({
+      id: Type.String({ description: 'Project ID (UUID)' }),
+    }),
+    async execute(_toolCallId, params) {
+      const { id } = params as { id: string };
+      const project = await brain.projects.findById(id);
+      return {
+        content: [
+          {
+            type: 'text' as const,
+            text: project ? JSON.stringify(project, null, 2) : `Project not found: ${id}`,
+          },
+        ],
+        details: undefined,
+      };
+    },
+  };
+
+  const listTasks: ToolDefinition = {
+    name: 'brain_list_tasks',
+    label: 'List Tasks',
+    description: 'List tasks, optionally filtered by project, mission, or status.',
+    parameters: Type.Object({
+      projectId: Type.Optional(Type.String({ description: 'Filter by project ID' })),
+      missionId: Type.Optional(Type.String({ description: 'Filter by mission ID' })),
+      status: Type.Optional(Type.String({ description: 'Filter by status' })),
+    }),
+    async execute(_toolCallId, params) {
+      const p = params as { projectId?: string; missionId?: string; status?: string };
+      type TaskStatus = 'not-started' | 'in-progress' | 'blocked' | 'done' | 'cancelled';
+      let tasks;
+      if (p.projectId) tasks = await brain.tasks.findByProject(p.projectId);
+      else if (p.missionId) tasks = await brain.tasks.findByMission(p.missionId);
+      else if (p.status) tasks = await brain.tasks.findByStatus(p.status as TaskStatus);
+      else tasks = await brain.tasks.findAll();
+      return {
+        content: [{ type: 'text' as const, text: JSON.stringify(tasks, null, 2) }],
+        details: undefined,
+      };
+    },
+  };
+
+  const createTask: ToolDefinition = {
+    name: 'brain_create_task',
+    label: 'Create Task',
+    description: 'Create a new task in the brain.',
+    parameters: Type.Object({
+      title: Type.String({ description: 'Task title' }),
+      description: Type.Optional(Type.String({ description: 'Task description' })),
+      projectId: Type.Optional(Type.String({ description: 'Project ID' })),
+      missionId: Type.Optional(Type.String({ description: 'Mission ID' })),
+      priority: Type.Optional(
+        Type.String({ description: 'Priority: low, medium, high, critical' }),
+      ),
+    }),
+    async execute(_toolCallId, params) {
+      const p = params as {
+        title: string;
+        description?: string;
+        projectId?: string;
+        missionId?: string;
+        priority?: string;
+      };
+      type Priority = 'low' | 'medium' | 'high' | 'critical';
+      const task = await brain.tasks.create({
+        ...p,
+        priority: p.priority as Priority | undefined,
+      });
+      return {
+        content: [{ type: 'text' as const, text: JSON.stringify(task, null, 2) }],
+        details: undefined,
+      };
+    },
+  };
+
+  const updateTask: ToolDefinition = {
+    name: 'brain_update_task',
+    label: 'Update Task',
+    description: 'Update an existing task.',
+    parameters: Type.Object({
+      id: Type.String({ description: 'Task ID' }),
+      title: Type.Optional(Type.String()),
+      description: Type.Optional(Type.String()),
+      status: Type.Optional(
+        Type.String({ description: 'not-started, in-progress, blocked, done, cancelled' }),
+      ),
+      priority: Type.Optional(Type.String()),
+    }),
+    async execute(_toolCallId, params) {
+      const { id, ...updates } = params as {
+        id: string;
+        title?: string;
+        description?: string;
+        status?: string;
+        priority?: string;
+      };
+      type TaskStatus = 'not-started' | 'in-progress' | 'blocked' | 'done' | 'cancelled';
+      type Priority = 'low' | 'medium' | 'high' | 'critical';
+      const task = await brain.tasks.update(id, {
+        ...updates,
+        status: updates.status as TaskStatus | undefined,
+        priority: updates.priority as Priority | undefined,
+      });
+      return {
+        content: [
+          {
+            type: 'text' as const,
+            text: task ? JSON.stringify(task, null, 2) : `Task not found: ${id}`,
+          },
+        ],
+        details: undefined,
+      };
+    },
+  };
+
+  const listMissions: ToolDefinition = {
+    name: 'brain_list_missions',
+    label: 'List Missions',
+    description: 'List all missions, optionally filtered by project.',
+    parameters: Type.Object({
+      projectId: Type.Optional(Type.String({ description: 'Filter by project ID' })),
+    }),
+    async execute(_toolCallId, params) {
+      const p = params as { projectId?: string };
+      const missions = p.projectId
+        ? await brain.missions.findByProject(p.projectId)
+        : await brain.missions.findAll();
+      return {
+        content: [{ type: 'text' as const, text: JSON.stringify(missions, null, 2) }],
+        details: undefined,
+      };
+    },
+  };
+
+  const listConversations: ToolDefinition = {
+    name: 'brain_list_conversations',
+    label: 'List Conversations',
+    description: 'List conversations for a user.',
+    parameters: Type.Object({
+      userId: Type.String({ description: 'User ID' }),
+    }),
+    async execute(_toolCallId, params) {
+      const { userId } = params as { userId: string };
+      const conversations = await brain.conversations.findAll(userId);
+      return {
+        content: [{ type: 'text' as const, text: JSON.stringify(conversations, null, 2) }],
+        details: undefined,
+      };
+    },
+  };
+
+  return [
+    listProjects,
+    getProject,
+    listTasks,
+    createTask,
+    updateTask,
+    listMissions,
+    listConversations,
+  ];
+}

apps/gateway/src/agent/tools/index.ts

@@ -0,0 +1 @@
+export { createBrainTools } from './brain-tools.js';

pnpm-lock.yaml

@@ -95,6 +95,9 @@ importers:
       '@opentelemetry/semantic-conventions':
         specifier: ^1.40.0
         version: 1.40.0
+      '@sinclair/typebox':
+        specifier: ^0.34.48
+        version: 0.34.48
       better-auth:
         specifier: ^1.5.5
         version: 1.5.5(drizzle-kit@0.31.9)(drizzle-orm@0.45.1(@opentelemetry/api@1.9.0)(@types/pg@8.15.6)(kysely@0.28.11)(postgres@3.4.8))(mongodb@7.1.0(socks@2.8.7))(next@16.1.6(@opentelemetry/api@1.9.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4))(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(vitest@2.1.9(@types/node@22.19.15))