@indemn/cli 1.0.0

package/src/mcp/tools.ts

@@ -0,0 +1,971 @@
+import { z } from 'zod/v3';
+import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import type { IndemnClient } from '../sdk/client.js';
+import { AgentsSDK } from '../sdk/agents.js';
+import { ConfigSDK } from '../sdk/config.js';
+import { FunctionsSDK } from '../sdk/functions.js';
+import { KBSDK } from '../sdk/kb.js';
+import { RubricsSDK } from '../sdk/rubrics.js';
+import { TestSetsSDK } from '../sdk/test-sets.js';
+import { EvalSDK } from '../sdk/eval.js';
+import { ChatSession } from '../sdk/chat.js';
+
+// Helper: format a successful MCP tool result
+function success(data: unknown) {
+  return { content: [{ type: 'text' as const, text: JSON.stringify(data, null, 2) }] };
+}
+
+// Helper: format an error MCP tool result
+function error(err: unknown) {
+  const message = err instanceof Error ? err.message : String(err);
+  return { content: [{ type: 'text' as const, text: `Error: ${message}` }], isError: true };
+}
+
+export function registerTools(server: McpServer, client: IndemnClient): void {
+  const agents = new AgentsSDK(client);
+  const config = new ConfigSDK(client);
+  const functions = new FunctionsSDK(client);
+  const kb = new KBSDK(client);
+  const rubrics = new RubricsSDK(client);
+  const testSets = new TestSetsSDK(client);
+  const evals = new EvalSDK(client);
+
+  // ---------------------------------------------------------------------------
+  // Agent tools
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'list_agents',
+    'List all agents in the organization',
+    { limit: z.number().optional().describe('Maximum number of agents to return') },
+    async (args) => {
+      try {
+        let result = await agents.listAgents();
+        if (args.limit) result = result.slice(0, args.limit);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_agent',
+    'Get details of a specific agent',
+    { agent_id: z.string().describe('The agent ID') },
+    async (args) => {
+      try {
+        const result = await agents.getAgent(args.agent_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'create_agent',
+    'Create a new agent',
+    {
+      name: z.string().describe('Agent name'),
+      channels: z.array(z.string()).optional().describe('Channels (e.g. ["WEB"])'),
+    },
+    async (args) => {
+      try {
+        const result = await agents.createAgent({ name: args.name, channels: args.channels });
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'update_agent',
+    'Update an existing agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      name: z.string().optional().describe('New agent name'),
+      description: z.string().optional().describe('New agent description'),
+      webhook_url: z.string().optional().describe('Webhook URL'),
+    },
+    async (args) => {
+      try {
+        const { agent_id, ...updates } = args;
+        const input: Record<string, unknown> = {};
+        if (updates.name !== undefined) input.name = updates.name;
+        if (updates.description !== undefined) input.description = updates.description;
+        if (updates.webhook_url !== undefined) input.webhook_url = updates.webhook_url;
+        const result = await agents.updateAgent(agent_id, input);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'delete_agent',
+    'Delete an agent',
+    { agent_id: z.string().describe('The agent ID to delete') },
+    async (args) => {
+      try {
+        await agents.deleteAgent(args.agent_id);
+        return success({ deleted: true, agent_id: args.agent_id });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'clone_agent',
+    'Clone an existing agent',
+    { agent_id: z.string().describe('The agent ID to clone') },
+    async (args) => {
+      try {
+        const result = await agents.cloneAgent(args.agent_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  // ---------------------------------------------------------------------------
+  // Config tools
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'get_config',
+    'Get the full configuration for an agent',
+    { agent_id: z.string().describe('The agent ID') },
+    async (args) => {
+      try {
+        const result = await config.getConfig(args.agent_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'set_config',
+    'Update configuration for an agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      system_prompt: z.string().optional().describe('System prompt text'),
+      first_message: z.string().optional().describe('First message text shown to users'),
+    },
+    async (args) => {
+      try {
+        const input: Record<string, unknown> = {};
+        if (args.system_prompt !== undefined) input.system_prompt = args.system_prompt;
+        if (args.first_message !== undefined) input.first_message = { text: args.first_message };
+        const result = await config.setConfig(args.agent_id, input);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_prompt',
+    'Get the system prompt for an agent',
+    { agent_id: z.string().describe('The agent ID') },
+    async (args) => {
+      try {
+        const result = await config.getPrompt(args.agent_id);
+        return success({ agent_id: args.agent_id, system_prompt: result ?? null });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'set_prompt',
+    'Set the system prompt for an agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      prompt: z.string().describe('The system prompt text'),
+    },
+    async (args) => {
+      try {
+        const result = await config.setPrompt(args.agent_id, args.prompt);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  // ---------------------------------------------------------------------------
+  // Function tools
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'list_functions',
+    'List all functions (tools) for an agent',
+    { agent_id: z.string().describe('The agent ID') },
+    async (args) => {
+      try {
+        const result = await functions.listFunctions(args.agent_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'create_function',
+    'Create a new function (tool) for an agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      type: z.string().describe('Function type (e.g. "rest_api")'),
+      description: z.string().optional().describe('What the function does'),
+      name: z.string().optional().describe('Function name'),
+      url: z.string().optional().describe('Endpoint URL'),
+      method: z.string().optional().describe('HTTP method (GET, POST, etc.)'),
+    },
+    async (args) => {
+      try {
+        const { agent_id, ...input } = args;
+        const result = await functions.createFunction(agent_id, input);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'update_function',
+    'Update an existing function (tool) for an agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      function_id: z.string().describe('The function ID'),
+      description: z.string().optional().describe('Updated description'),
+      name: z.string().optional().describe('Updated function name'),
+      url: z.string().optional().describe('Updated endpoint URL'),
+      method: z.string().optional().describe('Updated HTTP method'),
+    },
+    async (args) => {
+      try {
+        const { agent_id, function_id, ...input } = args;
+        const result = await functions.updateFunction(agent_id, function_id, input);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'delete_function',
+    'Delete a function (tool) from an agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      function_id: z.string().describe('The function ID to delete'),
+    },
+    async (args) => {
+      try {
+        await functions.deleteFunction(args.agent_id, args.function_id);
+        return success({ deleted: true, agent_id: args.agent_id, function_id: args.function_id });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'test_function',
+    'Test a function (tool) with optional parameters',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      function_id: z.string().describe('The function ID to test'),
+      params: z.record(z.unknown()).optional().describe('Test parameters to send'),
+    },
+    async (args) => {
+      try {
+        const result = await functions.testFunction(
+          args.agent_id,
+          args.function_id,
+          args.params as Record<string, unknown> | undefined,
+        );
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'export_functions',
+    'Export all functions from an agent',
+    { agent_id: z.string().describe('The agent ID') },
+    async (args) => {
+      try {
+        const result = await functions.exportFunctions(args.agent_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'import_functions',
+    'Import functions into an agent from a URL (e.g. OpenAPI spec)',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      url: z.string().describe('URL to import functions from'),
+    },
+    async (args) => {
+      try {
+        const result = await functions.importFunctions(args.agent_id, args.url);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'list_function_parameters',
+    'List parameters for a function',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      function_id: z.string().describe('The function ID'),
+    },
+    async (args) => {
+      try {
+        const result = await functions.listParameters(args.agent_id, args.function_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'add_function_parameter',
+    'Add a parameter to a function',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      function_id: z.string().describe('The function ID'),
+      name: z.string().describe('Parameter name'),
+      type: z.string().describe('Parameter type (string, number, boolean)'),
+      description: z.string().optional().describe('Parameter description'),
+      required: z.boolean().optional().describe('Whether the parameter is required'),
+    },
+    async (args) => {
+      try {
+        const { agent_id, function_id, ...param } = args;
+        const result = await functions.addParameter(agent_id, function_id, param);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'update_function_parameter',
+    'Update a function parameter',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      function_id: z.string().describe('The function ID'),
+      parameter_id: z.string().describe('The parameter ID'),
+      name: z.string().optional().describe('Updated name'),
+      type: z.string().optional().describe('Updated type'),
+      description: z.string().optional().describe('Updated description'),
+      required: z.boolean().optional().describe('Updated required flag'),
+    },
+    async (args) => {
+      try {
+        const { agent_id, function_id, parameter_id, ...param } = args;
+        const result = await functions.updateParameter(agent_id, function_id, parameter_id, param);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'delete_function_parameter',
+    'Delete a function parameter',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      function_id: z.string().describe('The function ID'),
+      parameter_id: z.string().describe('The parameter ID'),
+    },
+    async (args) => {
+      try {
+        await functions.deleteParameter(args.agent_id, args.function_id, args.parameter_id);
+        return success({ deleted: true });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'list_master_functions',
+    'List available master function templates',
+    {},
+    async () => {
+      try {
+        const result = await functions.listMasterFunctions();
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'list_llm_models',
+    'List available LLM models',
+    {},
+    async () => {
+      try {
+        const result = await functions.listLLMModels();
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  // ---------------------------------------------------------------------------
+  // Knowledge base tools
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'list_knowledge_bases',
+    'List all knowledge bases in the organization',
+    { limit: z.number().optional().describe('Maximum number of KBs to return') },
+    async (args) => {
+      try {
+        let result = await kb.listKBs();
+        if (args.limit) result = result.slice(0, args.limit);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_knowledge_base',
+    'Get details of a specific knowledge base',
+    { kb_id: z.string().describe('The knowledge base ID') },
+    async (args) => {
+      try {
+        const result = await kb.getKB(args.kb_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'create_knowledge_base',
+    'Create a new knowledge base',
+    {
+      name: z.string().describe('Knowledge base name'),
+      type: z.string().describe('Knowledge base type'),
+    },
+    async (args) => {
+      try {
+        const result = await kb.createKB({ name: args.name, type: args.type });
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'update_knowledge_base',
+    'Update an existing knowledge base',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      name: z.string().optional().describe('Updated name'),
+      description: z.string().optional().describe('Updated description'),
+    },
+    async (args) => {
+      try {
+        const input: Record<string, unknown> = {};
+        if (args.name !== undefined) input.name = args.name;
+        if (args.description !== undefined) input.description = args.description;
+        const result = await kb.updateKB(args.kb_id, input);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'delete_knowledge_base',
+    'Delete a knowledge base',
+    { kb_id: z.string().describe('The knowledge base ID to delete') },
+    async (args) => {
+      try {
+        await kb.deleteKB(args.kb_id);
+        return success({ deleted: true, kb_id: args.kb_id });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'list_kb_data',
+    'List data sources in a knowledge base',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      page: z.number().optional().describe('Page number'),
+      limit: z.number().optional().describe('Items per page'),
+    },
+    async (args) => {
+      try {
+        const query: Record<string, string> = {};
+        if (args.page !== undefined) query.page = String(args.page);
+        if (args.limit !== undefined) query.limit = String(args.limit);
+        const result = await kb.listDataSources(args.kb_id, Object.keys(query).length ? query : undefined);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'add_kb_data',
+    'Add a data source entry to a knowledge base. For URL scraping use scrape_url instead. For file uploads use upload_kb_file.',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      question: z.string().optional().describe('FAQ question'),
+      answer: z.string().optional().describe('FAQ answer'),
+      source_url: z.string().optional().describe('Source URL (initiates scraping via import endpoint)'),
+      content: z.string().optional().describe('Raw content'),
+    },
+    async (args) => {
+      try {
+        const { kb_id, ...input } = args;
+        if (input.source_url) {
+          const result = await kb.scrapeUrl(kb_id, input.source_url as string);
+          return success(result);
+        }
+        const result = await kb.addDataSource(kb_id, input);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'scrape_url',
+    'Scrape a website URL and add it as a data source to a knowledge base. Initiates async crawling.',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      url: z.string().describe('URL to scrape'),
+      crawl_mode: z.string().optional().describe('Crawl mode: "full" (entire site) or "selected"'),
+    },
+    async (args) => {
+      try {
+        const result = await kb.scrapeUrl(args.kb_id, args.url, args.crawl_mode);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'upload_kb_file',
+    'Upload a file to a knowledge base',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      file_path: z.string().describe('Path to the file to upload'),
+    },
+    async (args) => {
+      try {
+        const result = await kb.uploadFile(args.kb_id, args.file_path);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'update_kb_data',
+    'Update a data source entry in a knowledge base',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      source_id: z.string().describe('The data source ID'),
+      question: z.string().optional().describe('Updated question'),
+      answer: z.string().optional().describe('Updated answer'),
+      source_url: z.string().optional().describe('Updated source URL'),
+      content: z.string().optional().describe('Updated content'),
+    },
+    async (args) => {
+      try {
+        const { kb_id, source_id, ...input } = args;
+        const result = await kb.updateDataSource(kb_id, source_id, input);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'delete_kb_data',
+    'Delete a data source entry from a knowledge base',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      source_id: z.string().describe('The data source ID to delete'),
+    },
+    async (args) => {
+      try {
+        await kb.deleteDataSource(args.kb_id, args.source_id);
+        return success({ deleted: true, kb_id: args.kb_id, source_id: args.source_id });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'link_kb',
+    'Link knowledge bases to an agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      kb_ids: z.array(z.string()).describe('Array of knowledge base IDs to link'),
+    },
+    async (args) => {
+      try {
+        const result = await kb.linkKB(args.agent_id, args.kb_ids);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'unlink_kb',
+    'Unlink a knowledge base from an agent',
+    {
+      agent_id: z.string().describe('The agent ID'),
+      kb_id: z.string().describe('The knowledge base ID to unlink'),
+    },
+    async (args) => {
+      try {
+        await kb.unlinkKB(args.agent_id, args.kb_id);
+        return success({ unlinked: true, agent_id: args.agent_id, kb_id: args.kb_id });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'export_kb',
+    'Export knowledge base data',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      format: z.string().optional().describe('Export format (e.g. "json", "csv")'),
+    },
+    async (args) => {
+      try {
+        const result = await kb.exportKB(args.kb_id, args.format || 'json');
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'import_kb',
+    'Import data into a knowledge base',
+    {
+      kb_id: z.string().describe('The knowledge base ID'),
+      data: z.record(z.unknown()).describe('Data to import'),
+    },
+    async (args) => {
+      try {
+        const result = await kb.importKB(args.kb_id, args.data);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  // ---------------------------------------------------------------------------
+  // Chat tool
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'chat_with_agent',
+    'Send a message to an agent and get a response. Not interactive — sends one message and returns the reply.',
+    {
+      agent_id: z.string().describe('The agent ID (bot type) to chat with'),
+      message: z.string().describe('The message to send'),
+      session_id: z.string().optional().describe('Optional session ID to continue a conversation'),
+    },
+    async (args) => {
+      try {
+        const middlewareUrl = client.getHosts().middleware;
+        const chat = new ChatSession(args.agent_id, middlewareUrl, args.session_id);
+
+        // Track streaming state to distinguish greeting from response
+        let receiving = false;
+        let messageCount = 0;
+
+        chat.on('stream', (data: { isFirstChunk?: boolean; streamFinished?: boolean }) => {
+          if (data.isFirstChunk) receiving = true;
+          if (data.streamFinished) receiving = false;
+        });
+
+        const greetingMessages: string[] = [];
+        chat.on('message', (text: string) => {
+          greetingMessages.push(text);
+          messageCount++;
+        });
+
+        const sessionId = await chat.connect();
+
+        // Wait for greeting to finish streaming (up to 8s)
+        await new Promise<void>((resolve) => {
+          const check = setInterval(() => {
+            if (!receiving) { clearInterval(check); resolve(); }
+          }, 100);
+          setTimeout(() => { clearInterval(check); resolve(); }, 8000);
+        });
+
+        // Now swap the listener — next message is the actual response
+        const greetingCount = messageCount;
+        chat.removeAllListeners('message');
+
+        const response = await new Promise<string>((resolve, reject) => {
+          const timeout = setTimeout(() => {
+            chat.disconnect();
+            reject(new Error('Agent did not respond within 60 seconds'));
+          }, 60000);
+
+          chat.on('message', (text: string) => {
+            clearTimeout(timeout);
+            resolve(text);
+          });
+
+          chat.send(args.message);
+        });
+
+        chat.disconnect();
+
+        return success({ session_id: sessionId, greeting: greetingMessages.slice(0, greetingCount).join('\n'), response });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  // ---------------------------------------------------------------------------
+  // Eval tools
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'run_evaluation',
+    'Trigger an evaluation run for a bot against a test set',
+    {
+      bot_id: z.string().describe('The bot ID to evaluate'),
+      test_set_id: z.string().describe('The test set ID'),
+      rubric_id: z.string().optional().describe('Optional rubric ID'),
+      concurrency: z.number().optional().describe('Number of concurrent evaluations'),
+      eval_model: z.string().optional().describe('Model to use for evaluation judging'),
+      keep: z.boolean().optional().describe('Keep evaluation data after completion'),
+      limit: z.number().optional().describe('Limit items to evaluate'),
+      component_scope: z.string().optional().describe('Component scope filter'),
+      component_ids: z.array(z.string()).optional().describe('Component ID filter'),
+      mode: z.string().optional().describe('Evaluation mode'),
+    },
+    async (args) => {
+      try {
+        const input: Record<string, unknown> = {
+          bot_id: args.bot_id,
+          test_set_id: args.test_set_id,
+        };
+        if (args.rubric_id !== undefined) input.rubric_id = args.rubric_id;
+        if (args.concurrency !== undefined) input.concurrency = args.concurrency;
+        if (args.eval_model !== undefined) input.eval_model = args.eval_model;
+        if (args.keep !== undefined) input.keep = args.keep;
+        if (args.limit !== undefined) input.limit = args.limit;
+        if (args.component_scope !== undefined) input.component_scope = args.component_scope;
+        if (args.component_ids !== undefined) input.component_ids = args.component_ids;
+        if (args.mode !== undefined) input.mode = args.mode;
+        const result = await evals.triggerEvaluation(input as any);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_eval_status',
+    'Get the status of an evaluation run',
+    { run_id: z.string().describe('The evaluation run ID') },
+    async (args) => {
+      try {
+        const result = await evals.getRunStatus(args.run_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_eval_results',
+    'Get the detailed results of a completed evaluation run',
+    { run_id: z.string().describe('The evaluation run ID') },
+    async (args) => {
+      try {
+        const result = await evals.getRunResults(args.run_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_bot_context',
+    'Get the full context of a bot (prompt, tools, KBs, LLM config)',
+    { bot_id: z.string().describe('The bot ID') },
+    async (args) => {
+      try {
+        const result = await evals.getBotContext(args.bot_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  // ---------------------------------------------------------------------------
+  // Rubric tools
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'list_rubrics',
+    'List evaluation rubrics, optionally filtered by agent',
+    { agent_id: z.string().optional().describe('Filter rubrics by agent ID') },
+    async (args) => {
+      try {
+        const result = await rubrics.listRubrics(args.agent_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_rubric',
+    'Get details of a specific rubric',
+    { rubric_id: z.string().describe('The rubric ID') },
+    async (args) => {
+      try {
+        const result = await rubrics.getRubric(args.rubric_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'create_rubric',
+    'Create a new evaluation rubric',
+    {
+      name: z.string().describe('Rubric name'),
+      rules: z.array(z.record(z.unknown())).describe('Array of rubric rules'),
+      agent_id: z.string().optional().describe('Agent ID to associate the rubric with'),
+      description: z.string().optional().describe('Rubric description'),
+      severity_definitions: z.record(z.string()).optional().describe('Severity level definitions'),
+    },
+    async (args) => {
+      try {
+        const input: Record<string, unknown> = { name: args.name, rules: args.rules };
+        if (args.agent_id !== undefined) input.agent_id = args.agent_id;
+        if (args.description !== undefined) input.description = args.description;
+        if (args.severity_definitions !== undefined) input.severity_definitions = args.severity_definitions;
+        const result = await rubrics.createRubric(input as any);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'update_rubric',
+    'Update an existing rubric',
+    {
+      rubric_id: z.string().describe('The rubric ID'),
+      name: z.string().optional().describe('Updated name'),
+      description: z.string().optional().describe('Updated description'),
+      rules: z.array(z.record(z.unknown())).optional().describe('Updated rules'),
+    },
+    async (args) => {
+      try {
+        const input: Record<string, unknown> = {};
+        if (args.name !== undefined) input.name = args.name;
+        if (args.description !== undefined) input.description = args.description;
+        if (args.rules !== undefined) input.rules = args.rules;
+        const result = await rubrics.updateRubric(args.rubric_id, input as any);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'delete_rubric',
+    'Delete a rubric',
+    { rubric_id: z.string().describe('The rubric ID to delete') },
+    async (args) => {
+      try {
+        await rubrics.deleteRubric(args.rubric_id);
+        return success({ deleted: true, rubric_id: args.rubric_id });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'list_rubric_versions',
+    'List all versions of a rubric',
+    { rubric_id: z.string().describe('The rubric ID') },
+    async (args) => {
+      try {
+        const result = await rubrics.listVersions(args.rubric_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  // ---------------------------------------------------------------------------
+  // Test set tools
+  // ---------------------------------------------------------------------------
+
+  server.tool(
+    'list_test_sets',
+    'List test sets, optionally filtered by agent',
+    { agent_id: z.string().optional().describe('Filter test sets by agent ID') },
+    async (args) => {
+      try {
+        const result = await testSets.listTestSets(args.agent_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'get_test_set',
+    'Get details of a specific test set',
+    { test_set_id: z.string().describe('The test set ID') },
+    async (args) => {
+      try {
+        const result = await testSets.getTestSet(args.test_set_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'create_test_set',
+    'Create a new test set',
+    {
+      name: z.string().describe('Test set name'),
+      items: z.array(z.record(z.unknown())).describe('Array of test items'),
+      agent_id: z.string().optional().describe('Agent ID to associate the test set with'),
+      description: z.string().optional().describe('Test set description'),
+    },
+    async (args) => {
+      try {
+        const input: Record<string, unknown> = { name: args.name, items: args.items };
+        if (args.agent_id !== undefined) input.agent_id = args.agent_id;
+        if (args.description !== undefined) input.description = args.description;
+        const result = await testSets.createTestSet(input as any);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'update_test_set',
+    'Update an existing test set',
+    {
+      test_set_id: z.string().describe('The test set ID'),
+      name: z.string().optional().describe('Updated name'),
+      description: z.string().optional().describe('Updated description'),
+      items: z.array(z.record(z.unknown())).optional().describe('Updated test items'),
+    },
+    async (args) => {
+      try {
+        const input: Record<string, unknown> = {};
+        if (args.name !== undefined) input.name = args.name;
+        if (args.description !== undefined) input.description = args.description;
+        if (args.items !== undefined) input.items = args.items;
+        const result = await testSets.updateTestSet(args.test_set_id, input as any);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'delete_test_set',
+    'Delete a test set',
+    { test_set_id: z.string().describe('The test set ID to delete') },
+    async (args) => {
+      try {
+        await testSets.deleteTestSet(args.test_set_id);
+        return success({ deleted: true, test_set_id: args.test_set_id });
+      } catch (err) { return error(err); }
+    },
+  );
+
+  server.tool(
+    'list_test_set_versions',
+    'List all versions of a test set',
+    { test_set_id: z.string().describe('The test set ID') },
+    async (args) => {
+      try {
+        const result = await testSets.listVersions(args.test_set_id);
+        return success(result);
+      } catch (err) { return error(err); }
+    },
+  );
+}