| 1 | import { createSdkMcpServer, McpSdkServerConfigWithInstance } from '@anthropic-ai/claude-agent-sdk' |
| 2 | import { Injectable, Logger } from '@nestjs/common' |
| 3 | import { ContentGenerationTaskRepository } from '@yikart/mongodb' |
| 4 | import { Subject } from 'rxjs' |
| 5 | import { z } from 'zod' |
| 6 | import { AiAvailabilityService } from '../../ai-availability' |
| 7 | import { McpServerName } from '../agent.constants' |
| 8 | import { |
| 9 | AgentMessageType, |
| 10 | ContentGenerationTaskResultSchema, |
| 11 | ContentGenerationTaskTitleUpdatedChunkVo, |
| 12 | } from '../agent.vo' |
| 13 | import { successResult, wrapTool } from './mcp.utils' |
| 14 | |
| 15 | export { ContentGenerationTaskResultSchema } |
| 16 | |
| 17 | export enum TaskResultToolName { |
| 18 | OutputTaskResult = 'outputTaskResult', |
| 19 | SetTitle = 'setTitle', |
| 20 | } |
| 21 | |
| 22 | export enum UtilToolName { |
| 23 | GetCurrentTime = 'getCurrentTime', |
| 24 | Wait = 'wait', |
| 25 | OutputTaskResult = 'outputTaskResult', |
| 26 | SetTitle = 'setTitle', |
| 27 | } |
| 28 | |
| 29 | const getCurrentTimeSchema = z.object({}) |
| 30 | |
| 31 | const waitSchema = z.object({ |
| 32 | seconds: z.number().int().min(1).max(300), |
| 33 | }) |
| 34 | |
| 35 | @Injectable() |
| 36 | export class UtilMcp { |
| 37 | private readonly logger = new Logger(UtilMcp.name) |
| 38 | |
| 39 | constructor( |
| 40 | private readonly contentGenerateRepository: ContentGenerationTaskRepository, |
| 41 | private readonly aiAvailability: AiAvailabilityService, |
| 42 | ) { } |
| 43 | |
| 44 | getCurrentTime = wrapTool( |
| 45 | this.logger, |
| 46 | UtilToolName.GetCurrentTime, |
| 47 | 'Get the current server time in ISO 8601 format with timezone information. No parameters required. Returns current timestamp.', |
| 48 | getCurrentTimeSchema.shape, |
| 49 | async () => { |
| 50 | const now = new Date() |
| 51 | const isoString = now.toISOString() |
| 52 | return successResult(`Current time:\n- ISO 8601: ${isoString}`) |
| 53 | }, |
| 54 | this.aiAvailability, |
| 55 | ) |
| 56 | |
| 57 | wait = wrapTool( |
| 58 | this.logger, |
| 59 | UtilToolName.Wait, |
| 60 | 'Wait for a specified number of seconds before continuing. Use for controlling polling intervals for async tasks. Accepts seconds (1-300).', |
| 61 | waitSchema.shape, |
| 62 | async (args) => { |
| 63 | const seconds = Math.min(args.seconds, 300) |
| 64 | await new Promise(resolve => setTimeout(resolve, seconds * 1000)) |
| 65 | return successResult(`Waited for ${seconds} seconds`) |
| 66 | }, |
| 67 | this.aiAvailability, |
| 68 | ) |
| 69 | |
| 70 | createSetTitleTool( |
| 71 | taskId: string, |
| 72 | ) { |
| 73 | const titleUpdateSubject = new Subject<ContentGenerationTaskTitleUpdatedChunkVo>() |
| 74 | |
| 75 | const setTitleTool = wrapTool( |
| 76 | this.logger, |
| 77 | UtilToolName.SetTitle, |
| 78 | 'Set conversation title', |
| 79 | { |
| 80 | title: z.string().max(30), |
| 81 | }, |
| 82 | async (args) => { |
| 83 | const title = args.title |
| 84 | await this.contentGenerateRepository.updateById(taskId, { |
| 85 | title, |
| 86 | }) |
| 87 | this.logger.debug(`Title updated for task ${taskId}: ${title}`) |
| 88 | titleUpdateSubject.next(ContentGenerationTaskTitleUpdatedChunkVo.create({ |
| 89 | type: AgentMessageType.TitleUpdated, |
| 90 | taskId, |
| 91 | title, |
| 92 | })) |
| 93 | return successResult('Title updated successfully') |
| 94 | }, |
| 95 | this.aiAvailability, |
| 96 | ) |
| 97 | |
| 98 | return [ |
| 99 | setTitleTool, |
| 100 | titleUpdateSubject.asObservable(), |
| 101 | () => titleUpdateSubject.complete(), |
| 102 | ] as const |
| 103 | } |
| 104 | |
| 105 | readonly server: McpSdkServerConfigWithInstance = createSdkMcpServer({ |
| 106 | name: McpServerName.Util, |
| 107 | version: '1.0.0', |
| 108 | tools: [ |
| 109 | this.wait, |
| 110 | this.getCurrentTime, |
| 111 | ], |
| 112 | }) |
| 113 | } |
| 114 |