package mcphost import ( "context" "fmt" "os" "path/filepath" "strings" "time" "github.com/bytedance/sonic" "github.com/mark3labs/mcp-go/mcp" "github.com/rs/zerolog/log" "github.com/httprunner/httprunner/v5/uixt" "github.com/httprunner/httprunner/v5/uixt/option" ) // MCPToolRecord represents a single tool record in the database // Each record contains detailed information about a tool and its server type MCPToolRecord struct { ToolID string `json:"tool_id"` // Unique identifier for the tool record BizID string `json:"biz_id"` // Business ID of the tool VisibleRange int `json:"visible_range"` // Visible range of the tool, 0: visible to biz, 1: visible to all ToolType string `json:"tool_type"` // Type of the tool ServerName string `json:"mcp_server"` // Name of the MCP server ToolName string `json:"tool_name"` // Name of the tool Description string `json:"description"` // Tool description Parameters string `json:"parameters"` // Tool input parameters in JSON format Returns string `json:"return_desc"` // Tool return value format in JSON format TeardownPair string `json:"teardown_pair"` // Teardown pair of the tool Examples string `json:"examples"` // Examples of the tool SupportPatterns string `json:"support_patterns"` // Support pattern of the tool CreatedAt time.Time `json:"created_at"` // Record creation time LastUpdatedAt time.Time `json:"last_updated_at"` // Record last update time } // DocStringInfo contains the parsed information from a Python docstring type DocStringInfo struct { Description string Parameters map[string]string Returns map[string]string } // extractDocStringInfo extracts information from a Python docstring // Example input: // """Get weather alerts for a US state. // // Args: // state: Two-letter US state code (e.g. CA, NY) // // Returns: // alerts: List of active weather alerts for the specified state // error: Error message if the request fails // """ func extractDocStringInfo(docstring string) DocStringInfo { info := DocStringInfo{ Parameters: make(map[string]string), Returns: make(map[string]string), } // Find the Args and Returns sections argsIndex := strings.Index(docstring, "Args:") returnsIndex := strings.Index(docstring, "Returns:") // Extract description (everything before Args) if argsIndex != -1 { info.Description = strings.TrimSpace(docstring[:argsIndex]) } else if returnsIndex != -1 { info.Description = strings.TrimSpace(docstring[:returnsIndex]) } else { info.Description = strings.TrimSpace(docstring) return info } // Helper function to extract key-value pairs from a section extractSection := func(content string) map[string]string { result := make(map[string]string) lines := strings.Split(content, "\n") for _, line := range lines { line = strings.TrimSpace(line) if line == "" { continue } parts := strings.SplitN(line, ":", 2) if len(parts) != 2 { continue } key := strings.TrimSpace(parts[0]) value := strings.TrimSpace(parts[1]) if key != "" && value != "" { result[key] = value } } return result } // Extract Args section if argsIndex != -1 { endIndex := returnsIndex if endIndex == -1 { endIndex = len(docstring) } argsContent := docstring[argsIndex+len("Args:") : endIndex] info.Parameters = extractSection(argsContent) } // Extract Returns section if returnsIndex != -1 { returnsContent := docstring[returnsIndex+len("Returns:"):] info.Returns = extractSection(returnsContent) } return info } // ActionToolProvider defines the interface for MCP servers that provide ActionTool implementations type ActionToolProvider interface { GetToolByAction(actionName option.ActionName) uixt.ActionTool } // ConvertToolsToRecords converts []MCPTools to a list of database records func (host *MCPHost) ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord { var records []MCPToolRecord now := time.Now() for _, mcpTools := range tools { if mcpTools.Err != nil { log.Error().Str("server", mcpTools.ServerName).Err(mcpTools.Err).Msg("skip tools conversion due to error") continue } for _, tool := range mcpTools.Tools { record := host.convertSingleToolToRecord(mcpTools.ServerName, tool, now) records = append(records, record) } } return records } // convertSingleToolToRecord converts a single MCP tool to a database record func (host *MCPHost) convertSingleToolToRecord(serverName string, tool mcp.Tool, timestamp time.Time) MCPToolRecord { // Generate unique ID id := fmt.Sprintf("%s__%s", serverName, tool.Name) // Extract description from docstring info := extractDocStringInfo(tool.Description) // Extract parameters paramsJSON := host.extractParameters(tool, info) // Extract returns returnsJSON := host.extractReturns(serverName, tool.Name, info) return MCPToolRecord{ ToolID: id, VisibleRange: 1, ToolType: "Hrp", ServerName: serverName, ToolName: tool.Name, Description: info.Description, Parameters: paramsJSON, Returns: returnsJSON, CreatedAt: timestamp, LastUpdatedAt: timestamp, } } // extractParameters extracts parameter information from tool schema or docstring func (host *MCPHost) extractParameters(tool mcp.Tool, info DocStringInfo) string { // Priority 1: Extract from InputSchema.Properties if len(tool.InputSchema.Properties) > 0 { return host.extractParametersFromSchema(tool.InputSchema.Properties) } // Priority 2: Extract from docstring if len(info.Parameters) > 0 { return host.marshalToJSON(info.Parameters, "docstring parameters") } return "{}" } // extractParametersFromSchema extracts parameters from MCP tool input schema func (host *MCPHost) extractParametersFromSchema(properties map[string]interface{}) string { schemaParams := make(map[string]string) for propName, propValue := range properties { propMap, ok := propValue.(map[string]interface{}) if !ok { continue } description := host.getPropertyDescription(propMap) schemaParams[propName] = description } return host.marshalToJSON(schemaParams, "schema parameters") } // getPropertyDescription extracts description from property map func (host *MCPHost) getPropertyDescription(propMap map[string]interface{}) string { if desc, exists := propMap["description"]; exists { if descStr, ok := desc.(string); ok { return descStr } } // Fallback to type information if propType, exists := propMap["type"]; exists { if typeStr, ok := propType.(string); ok { return fmt.Sprintf("Parameter of type %s", typeStr) } } return "Parameter" } // extractReturns extracts return value information from ActionTool or docstring func (host *MCPHost) extractReturns(serverName, toolName string, info DocStringInfo) string { // Priority 1: Get from ActionTool interface if available if actionToolProvider := host.getActionToolProvider(serverName); actionToolProvider != nil { if actionTool := actionToolProvider.GetToolByAction(option.ActionName(toolName)); actionTool != nil { returnSchema := uixt.GenerateReturnSchema(actionTool) if len(returnSchema) > 0 { return host.marshalToJSON(returnSchema, "return schema") } } } // Priority 2: Use docstring returns as fallback if len(info.Returns) > 0 { return host.marshalToJSON(info.Returns, "docstring returns") } return "{}" } // marshalToJSON marshals data to JSON string with error handling func (host *MCPHost) marshalToJSON(data interface{}, dataType string) string { jsonBytes, err := sonic.MarshalString(data) if err != nil { log.Warn().Interface("data", data).Err(err). Msgf("failed to marshal %s to JSON", dataType) return "{}" } return jsonBytes } // ExportToolsToJSON dumps MCP tools to JSON file func (h *MCPHost) ExportToolsToJSON(ctx context.Context, dumpPath string) error { // get all tools tools := h.GetTools(ctx) // convert to records records := h.ConvertToolsToRecords(tools) // convert to JSON recordsJSON, err := sonic.MarshalIndent(records, "", " ") if err != nil { return fmt.Errorf("failed to marshal records to JSON: %w", err) } // create output directory outputDir := filepath.Dir(dumpPath) if outputDir != "." { if err := os.MkdirAll(outputDir, 0o754); err != nil { return fmt.Errorf("failed to create output directory: %w", err) } } // write to file if err := os.WriteFile(dumpPath, []byte(recordsJSON), 0o644); err != nil { return fmt.Errorf("failed to write records to file: %w", err) } log.Info().Str("path", dumpPath).Msg("Tools records exported successfully") return nil }