mirror of
https://github.com/httprunner/httprunner.git
synced 2026-05-11 10:00:23 +08:00
feat: hrp mcphost
This commit is contained in:
lilong.129
@@ -99,6 +99,7 @@ func LoadMCPConfig(configPath string) (*MCPConfig, error) {
|
||||
}
|
||||
configPath = filepath.Join(homeDir, ".mcp.json")
|
||||
}
|
||||
configPath = os.ExpandEnv(configPath)
|
||||
|
||||
// Check if config file exists
|
||||
if _, err := os.Stat(configPath); os.IsNotExist(err) {
|
||||
|
||||
185
pkg/mcphost/dump.go
Normal file
185
pkg/mcphost/dump.go
Normal file
@@ -0,0 +1,185 @@
|
||||
package mcphost
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/bytedance/sonic"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// ConvertToolsToRecords converts map[string]MCPTools to a list of database records
|
||||
func ConvertToolsToRecords(toolsMap map[string]MCPTools) []MCPToolRecord {
|
||||
var records []MCPToolRecord
|
||||
now := time.Now()
|
||||
|
||||
for serverName, mcpTools := range toolsMap {
|
||||
if mcpTools.Err != nil {
|
||||
log.Error().Str("server", serverName).Err(mcpTools.Err).Msg("skip tools conversion due to error")
|
||||
continue
|
||||
}
|
||||
|
||||
for _, tool := range mcpTools.Tools {
|
||||
// Generate unique ID by combining server name and tool name
|
||||
id := fmt.Sprintf("%s_%s", 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: serverName,
|
||||
ToolName: tool.Name,
|
||||
Description: info.Description,
|
||||
Parameters: paramsJSON,
|
||||
Returns: returnsJSON,
|
||||
CreatedAt: now,
|
||||
LastUpdatedAt: now,
|
||||
}
|
||||
|
||||
records = append(records, record)
|
||||
}
|
||||
}
|
||||
|
||||
return records
|
||||
}
|
||||
|
||||
// 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)
|
||||
// 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, 0o755); 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
|
||||
}
|
||||
241
pkg/mcphost/dump_test.go
Normal file
241
pkg/mcphost/dump_test.go
Normal file
@@ -0,0 +1,241 @@
|
||||
package mcphost
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestConvertToolsToRecordsFromFile(t *testing.T) {
|
||||
hub, err := NewMCPHost("./testdata/test.mcp.json")
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx := context.Background()
|
||||
err = hub.InitServers(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
// use ExportToolsToJSON to dump tools to JSON file
|
||||
err = hub.ExportToolsToJSON(ctx, "./tools_records.json")
|
||||
require.NoError(t, err)
|
||||
|
||||
// read the exported JSON file
|
||||
data, err := os.ReadFile("./tools_records.json")
|
||||
require.NoError(t, err)
|
||||
|
||||
// parse the exported JSON data
|
||||
var records []MCPToolRecord
|
||||
err = json.Unmarshal(data, &records)
|
||||
require.NoError(t, err)
|
||||
|
||||
// verify the number of records
|
||||
assert.NotEmpty(t, records, "Exported records should not be empty")
|
||||
|
||||
t.Logf("Tools records written to ./tools_records.json")
|
||||
}
|
||||
|
||||
func TestExtractDocStringInfo(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
docstring string
|
||||
want DocStringInfo
|
||||
}{
|
||||
{
|
||||
name: "complete docstring with args and returns",
|
||||
docstring: `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
|
||||
`,
|
||||
want: DocStringInfo{
|
||||
Description: "Get weather alerts for a US state.",
|
||||
Parameters: map[string]string{
|
||||
"state": "Two-letter US state code (e.g. CA, NY)",
|
||||
},
|
||||
Returns: map[string]string{
|
||||
"alerts": "List of active weather alerts for the specified state",
|
||||
"error": "Error message if the request fails",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "docstring with only args",
|
||||
docstring: `Do screen swipe action.
|
||||
|
||||
Args:
|
||||
direction: swipe direction (up, down)
|
||||
`,
|
||||
want: DocStringInfo{
|
||||
Description: "Do screen swipe action.",
|
||||
Parameters: map[string]string{
|
||||
"direction": "swipe direction (up, down)",
|
||||
},
|
||||
Returns: map[string]string{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "docstring with only description",
|
||||
docstring: "Simple tool with no parameters.",
|
||||
want: DocStringInfo{
|
||||
Description: "Simple tool with no parameters.",
|
||||
Parameters: map[string]string{},
|
||||
Returns: map[string]string{},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "docstring with multiple parameters",
|
||||
docstring: `Perform complex operation.
|
||||
|
||||
Args:
|
||||
param1: first parameter description
|
||||
param2: second parameter description
|
||||
param3: third parameter description
|
||||
|
||||
Returns:
|
||||
result: operation result
|
||||
`,
|
||||
want: DocStringInfo{
|
||||
Description: "Perform complex operation.",
|
||||
Parameters: map[string]string{
|
||||
"param1": "first parameter description",
|
||||
"param2": "second parameter description",
|
||||
"param3": "third parameter description",
|
||||
},
|
||||
Returns: map[string]string{
|
||||
"result": "operation result",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := extractDocStringInfo(tt.docstring)
|
||||
assert.Equal(t, tt.want.Description, got.Description)
|
||||
assert.Equal(t, tt.want.Parameters, got.Parameters)
|
||||
assert.Equal(t, tt.want.Returns, got.Returns)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertToolsToRecords(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
toolsMap map[string]MCPTools
|
||||
want []MCPToolRecord
|
||||
}{
|
||||
{
|
||||
name: "convert weather tool",
|
||||
toolsMap: map[string]MCPTools{
|
||||
"weather": {
|
||||
Name: "weather",
|
||||
Tools: []mcp.Tool{
|
||||
{
|
||||
Name: "get_alerts",
|
||||
Description: `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
|
||||
`,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
want: []MCPToolRecord{
|
||||
{
|
||||
ToolID: "weather_get_alerts",
|
||||
ServerName: "weather",
|
||||
ToolName: "get_alerts",
|
||||
Description: "Get weather alerts for a US state.",
|
||||
Parameters: `{"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"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "convert multiple tools",
|
||||
toolsMap: map[string]MCPTools{
|
||||
"ui": {
|
||||
Name: "ui",
|
||||
Tools: []mcp.Tool{
|
||||
{
|
||||
Name: "swipe",
|
||||
Description: `Do screen swipe action.
|
||||
|
||||
Args:
|
||||
direction: swipe direction (up, down)
|
||||
`,
|
||||
},
|
||||
{
|
||||
Name: "tap",
|
||||
Description: "Tap on screen at specified position.",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
want: []MCPToolRecord{
|
||||
{
|
||||
ToolID: "ui_swipe",
|
||||
ServerName: "ui",
|
||||
ToolName: "swipe",
|
||||
Description: "Do screen swipe action.",
|
||||
Parameters: `{"direction":"swipe direction (up, down)"}`,
|
||||
Returns: "{}",
|
||||
},
|
||||
{
|
||||
ToolID: "ui_tap",
|
||||
ServerName: "ui",
|
||||
ToolName: "tap",
|
||||
Description: "Tap on screen at specified position.",
|
||||
Parameters: "{}",
|
||||
Returns: "{}",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ConvertToolsToRecords(tt.toolsMap)
|
||||
|
||||
// Compare each record
|
||||
require.Equal(t, len(tt.want), len(got))
|
||||
for i := range tt.want {
|
||||
assert.Equal(t, tt.want[i].ToolID, got[i].ToolID)
|
||||
assert.Equal(t, tt.want[i].ServerName, got[i].ServerName)
|
||||
assert.Equal(t, tt.want[i].ToolName, got[i].ToolName)
|
||||
assert.Equal(t, tt.want[i].Description, got[i].Description)
|
||||
|
||||
// Compare JSON content (ignoring whitespace differences)
|
||||
var wantParams, gotParams, wantReturns, gotReturns map[string]string
|
||||
require.NoError(t, json.Unmarshal([]byte(tt.want[i].Parameters), &wantParams))
|
||||
require.NoError(t, json.Unmarshal([]byte(got[i].Parameters), &gotParams))
|
||||
require.NoError(t, json.Unmarshal([]byte(tt.want[i].Returns), &wantReturns))
|
||||
require.NoError(t, json.Unmarshal([]byte(got[i].Returns), &gotReturns))
|
||||
|
||||
assert.Equal(t, wantParams, gotParams)
|
||||
assert.Equal(t, wantReturns, gotReturns)
|
||||
|
||||
// Verify timestamps are recent (within last 5 seconds)
|
||||
now := time.Now()
|
||||
assert.True(t, now.Sub(got[i].CreatedAt) < 5*time.Second, "CreatedAt should be recent")
|
||||
assert.True(t, now.Sub(got[i].LastUpdatedAt) < 5*time.Second, "LastUpdatedAt should be recent")
|
||||
// CreatedAt and LastUpdatedAt should be the same
|
||||
assert.Equal(t, got[i].CreatedAt, got[i].LastUpdatedAt)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -169,9 +169,9 @@ func TestConcurrentOperations(t *testing.T) {
|
||||
|
||||
// Test concurrent tool invocations
|
||||
done := make(chan bool)
|
||||
timeout := time.After(10 * time.Second) // Increase timeout to 10 seconds
|
||||
timeout := time.After(30 * time.Second) // Increase timeout to 30 seconds
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
for i := 0; i < 3; i++ { // Reduce number of concurrent operations to 3
|
||||
go func() {
|
||||
result, err := host.InvokeTool(ctx, "weather", "get_alerts",
|
||||
map[string]interface{}{"state": "CA"},
|
||||
@@ -183,7 +183,7 @@ func TestConcurrentOperations(t *testing.T) {
|
||||
}
|
||||
|
||||
// Wait for all goroutines to complete
|
||||
for i := 0; i < 5; i++ {
|
||||
for i := 0; i < 3; i++ { // Update loop count to match the number of goroutines
|
||||
select {
|
||||
case <-done:
|
||||
// Success
|
||||
|
||||
Reference in New Issue
Block a user