GithubBackup/httprunner
1
0
Fork 0
mirror of https://github.com/httprunner/httprunner.git synced 2026-07-11 15:31:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 重构 MCP 工具导出逻辑并完善返回值类型系统

Browse Source
This commit is contained in:
lilong.129
2025-05-31 00:28:24 +08:00
parent 2a392f204b
commit 9089bd9324
7 changed files with 924 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

175
mcphost/dump.go
View File

@@ -9,20 +9,30 @@ import (
"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
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:"returns"` // Tool return value format in JSON format
CreatedAt time.Time `json:"created_at"` // Record creation time
LastUpdatedAt time.Time `json:"last_updated_at"` // Record last update time
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
@@ -109,8 +119,13 @@ func extractDocStringInfo(docstring string) DocStringInfo {
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 ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
func (host *MCPHost) ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
var records []MCPToolRecord
now := time.Now()
@@ -121,36 +136,7 @@ func ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
}
for _, tool := range mcpTools.Tools {
// Generate unique ID by combining server name and tool name
id := fmt.Sprintf("%s__%s", mcpTools.ServerName, tool.Name)
// Extract docstring information
info := extractDocStringInfo(tool.Description)
// Convert parameters and returns to JSON
paramsJSON, err := sonic.MarshalString(info.Parameters)
if err != nil {
log.Warn().Interface("params", info.Parameters).Err(err).Msg("failed to marshal parameters to JSON")
paramsJSON = "{}"
}
returnsJSON, err := sonic.MarshalString(info.Returns)
if err != nil {
log.Warn().Interface("returns", info.Returns).Err(err).Msg("failed to marshal returns to JSON")
returnsJSON = "{}"
}
record := MCPToolRecord{
ToolID: id,
ServerName: mcpTools.ServerName,
ToolName: tool.Name,
Description: info.Description,
Parameters: paramsJSON,
Returns: returnsJSON,
CreatedAt: now,
LastUpdatedAt: now,
}
record := host.convertSingleToolToRecord(mcpTools.ServerName, tool, now)
records = append(records, record)
}
}
@@ -158,12 +144,121 @@ func ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
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: "edge",
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 := actionTool.ReturnSchema()
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 := ConvertToolsToRecords(tools)
records := h.ConvertToolsToRecords(tools)
// convert to JSON
recordsJSON, err := sonic.MarshalIndent(records, "", " ")
if err != nil {

225
mcphost/dump_test.go
View File

@@ -124,6 +124,11 @@ func TestExtractDocStringInfo(t *testing.T) {
}
func TestConvertToolsToRecords(t *testing.T) {
// Create a mock MCPHost for testing
host := &MCPHost{
connections: make(map[string]*Connection),
}
tests := []struct {
name string
tools []MCPTools
@@ -152,7 +157,7 @@ func TestConvertToolsToRecords(t *testing.T) {
},
want: []MCPToolRecord{
{
ToolID: "weather_get_alerts",
ToolID: "weather__get_alerts",
ServerName: "weather",
ToolName: "get_alerts",
Description: "Get weather alerts for a US state.",
@@ -184,7 +189,7 @@ func TestConvertToolsToRecords(t *testing.T) {
},
want: []MCPToolRecord{
{
ToolID: "ui_swipe",
ToolID: "ui__swipe",
ServerName: "ui",
ToolName: "swipe",
Description: "Do screen swipe action.",
@@ -192,7 +197,7 @@ func TestConvertToolsToRecords(t *testing.T) {
Returns: "{}",
},
{
ToolID: "ui_tap",
ToolID: "ui__tap",
ServerName: "ui",
ToolName: "tap",
Description: "Tap on screen at specified position.",
@@ -201,11 +206,47 @@ func TestConvertToolsToRecords(t *testing.T) {
},
},
},
{
name: "convert tool with InputSchema",
tools: []MCPTools{
{
ServerName: "test",
Tools: []mcp.Tool{
{
Name: "test_tool",
Description: "Test tool with input schema",
InputSchema: mcp.ToolInputSchema{
Type: "object",
Properties: map[string]interface{}{
"param1": map[string]interface{}{
"type": "string",
"description": "First parameter",
},
"param2": map[string]interface{}{
"type": "number",
},
},
},
},
},
},
},
want: []MCPToolRecord{
{
ToolID: "test__test_tool",
ServerName: "test",
ToolName: "test_tool",
Description: "Test tool with input schema",
Parameters: `{"param1":"First parameter","param2":"Parameter of type number"}`,
Returns: "{}",
},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := ConvertToolsToRecords(tt.tools)
got := host.ConvertToolsToRecords(tt.tools)
// Compare each record
require.Equal(t, len(tt.want), len(got))
@@ -235,3 +276,179 @@ func TestConvertToolsToRecords(t *testing.T) {
})
}
}
// TestExtractParameters tests the extractParameters method
func TestExtractParameters(t *testing.T) {
host := &MCPHost{}
tests := []struct {
name string
tool mcp.Tool
info DocStringInfo
expected string
}{
{
name: "extract from InputSchema",
tool: mcp.Tool{
InputSchema: mcp.ToolInputSchema{
Properties: map[string]interface{}{
"param1": map[string]interface{}{
"type": "string",
"description": "First parameter",
},
"param2": map[string]interface{}{
"type": "number",
},
},
},
},
info: DocStringInfo{Parameters: map[string]string{"old": "old param"}},
expected: `{"param1":"First parameter","param2":"Parameter of type number"}`,
},
{
name: "fallback to docstring",
tool: mcp.Tool{},
info: DocStringInfo{
Parameters: map[string]string{
"param": "parameter description",
},
},
expected: `{"param":"parameter description"}`,
},
{
name: "empty parameters",
tool: mcp.Tool{},
info: DocStringInfo{},
expected: "{}",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := host.extractParameters(tt.tool, tt.info)
assert.Equal(t, tt.expected, got)
})
}
}
// TestExtractReturns tests the extractReturns method
func TestExtractReturns(t *testing.T) {
host := &MCPHost{
connections: make(map[string]*Connection),
}
tests := []struct {
name string
serverName string
toolName string
info DocStringInfo
expected string
}{
{
name: "fallback to docstring returns",
serverName: "unknown_server",
toolName: "unknown_tool",
info: DocStringInfo{
Returns: map[string]string{
"result": "operation result",
"error": "error message",
},
},
expected: `{"error":"error message","result":"operation result"}`,
},
{
name: "empty returns",
serverName: "unknown_server",
toolName: "unknown_tool",
info: DocStringInfo{},
expected: "{}",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := host.extractReturns(tt.serverName, tt.toolName, tt.info)
assert.Equal(t, tt.expected, got)
})
}
}
// TestGetPropertyDescription tests the getPropertyDescription method
func TestGetPropertyDescription(t *testing.T) {
host := &MCPHost{}
tests := []struct {
name string
propMap map[string]interface{}
expected string
}{
{
name: "with description",
propMap: map[string]interface{}{
"type": "string",
"description": "Parameter description",
},
expected: "Parameter description",
},
{
name: "without description, with type",
propMap: map[string]interface{}{
"type": "number",
},
expected: "Parameter of type number",
},
{
name: "without description and type",
propMap: map[string]interface{}{},
expected: "Parameter",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := host.getPropertyDescription(tt.propMap)
assert.Equal(t, tt.expected, got)
})
}
}
// TestMarshalToJSON tests the marshalToJSON method
func TestMarshalToJSON(t *testing.T) {
host := &MCPHost{}
tests := []struct {
name string
data interface{}
dataType string
expected string
}{
{
name: "valid map",
data: map[string]string{
"key1": "value1",
"key2": "value2",
},
dataType: "test data",
expected: `{"key1":"value1","key2":"value2"}`,
},
{
name: "empty map",
data: map[string]string{},
dataType: "test data",
expected: "{}",
},
{
name: "invalid data (channel)",
data: make(chan int),
dataType: "test data",
expected: "{}",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := host.marshalToJSON(tt.data, tt.dataType)
assert.Equal(t, tt.expected, got)
})
}
}

15
mcphost/host.go
View File

@@ -558,3 +558,18 @@ func (h *MCPHost) forceCloseAll() {
delete(h.connections, name)
}
}
// getActionToolProvider returns an ActionToolProvider for the given server name if available
// This method checks if the MCP server implements the ActionToolProvider interface
func (h *MCPHost) getActionToolProvider(serverName string) ActionToolProvider {
h.mu.RLock()
defer h.mu.RUnlock()
if conn, exists := h.connections[serverName]; exists {
// Check if the client directly implements ActionToolProvider interface
if actionToolProvider, ok := conn.Client.(ActionToolProvider); ok {
return actionToolProvider
}
}
return nil
}