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

Register brain data layer as LLM-callable tools in Pi SDK agent
sessions. Agents can now query/create/update projects, tasks,
missions, and conversations through tool calls.

- 7 brain tools: list/get projects, list/create/update tasks,
  list missions, list conversations
- Tools registered via customTools in createAgentSession()
- Uses TypeBox schemas for parameter validation

Closes #22

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-12 22:18:09 -05:00
parent 7485f32e69
commit d0d3683651
5 changed files with 204 additions and 2 deletions

View File

@@ -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",

View File

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

View File

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

View File

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

3
pnpm-lock.yaml generated
View File

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