LLMGuardX Dev
d534865b33
- 节点任务:管理端向节点池下发(Prompt/HTTP/MCP/Skill/数字员工/自定义), 指定节点或池路由,认领 SKIP LOCKED + 15 分钟租约,认领令牌防重放上报, 失败 30s×次数退避重入队,达上限 failed,支持取消/重试,完成与失败站内信。 - 个人智能体安全策略:auto_approve_tools 跳过个人调用审批门; rate_limit_multiplier 按 (tool,user) 独立窗口放宽个人限流(全局额度不受影响)。 - 修复存量缺陷:/v1/agent/nodes/ 未挂 publicMux,节点心跳/认领端点在部署 拓扑下不可达。 - 迁移 000046;任务全链路集成测试连真实库通过,HTTP 端到端验证 (下发→认领→伪造令牌拒绝→上报→succeeded,列表不泄露认领令牌); 25 包测试通过,前后端构建通过。
534 lines
20 KiB
Go
534 lines
20 KiB
Go
package workbench
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"aigateway.local/core/internal/platform/cryptox"
|
|
"aigateway.local/core/internal/provider"
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
type ToolService struct {
|
|
assets *Service
|
|
cipher cryptox.Cipher
|
|
allowPrivate bool
|
|
}
|
|
|
|
func NewToolService(assets *Service, cipher cryptox.Cipher, allowPrivate bool) *ToolService {
|
|
return &ToolService{assets: assets, cipher: cipher, allowPrivate: allowPrivate}
|
|
}
|
|
|
|
func (s *ToolService) validate(ctx context.Context, input *ToolInput, create bool) error {
|
|
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
|
|
input.Name = strings.TrimSpace(input.Name)
|
|
input.Description = strings.TrimSpace(input.Description)
|
|
input.EndpointURL = strings.TrimSpace(input.EndpointURL)
|
|
input.HTTPMethod = strings.ToUpper(strings.TrimSpace(input.HTTPMethod))
|
|
if !codeRE.MatchString(input.Code) || input.Name == "" || len(input.Name) > 128 || len(input.Description) > 4000 {
|
|
return errors.New("工具编码、名称或描述格式无效")
|
|
}
|
|
if input.HTTPMethod == "" {
|
|
input.HTTPMethod = "POST"
|
|
}
|
|
switch input.HTTPMethod {
|
|
case "GET", "POST", "PUT", "PATCH", "DELETE":
|
|
default:
|
|
return errors.New("不支持的 HTTP 方法")
|
|
}
|
|
validated, err := provider.ValidateBaseURL(ctx, input.EndpointURL, s.allowPrivate)
|
|
if err != nil {
|
|
return fmt.Errorf("工具端点校验失败: %w", err)
|
|
}
|
|
input.EndpointURL = validated
|
|
if input.TimeoutSeconds == 0 {
|
|
input.TimeoutSeconds = 15
|
|
}
|
|
if input.TimeoutSeconds < 1 || input.TimeoutSeconds > 120 {
|
|
return errors.New("超时应在 1-120 秒之间")
|
|
}
|
|
if input.RateLimitRPM < 0 || input.RateLimitRPM > 100000 {
|
|
return errors.New("工具限流应在 0-100000 RPM 之间")
|
|
}
|
|
input.DepartmentIDs, err = normalizeStrings(input.DepartmentIDs, 100)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(input.InputSchema) == 0 {
|
|
input.InputSchema = json.RawMessage(`{}`)
|
|
}
|
|
var schema map[string]any
|
|
if err = json.Unmarshal(input.InputSchema, &schema); err != nil {
|
|
return errors.New("input_schema 必须是 JSON 对象")
|
|
}
|
|
normalized, err := json.Marshal(schema)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
input.InputSchema = normalized
|
|
if create && input.Headers == nil {
|
|
input.Headers = map[string]string{}
|
|
}
|
|
for key, value := range input.Headers {
|
|
if strings.TrimSpace(key) == "" || len(key) > 128 || strings.ContainsAny(key, "\r\n") || len(value) > 8192 || strings.ContainsAny(value, "\r\n") {
|
|
return errors.New("工具请求头格式无效")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
const toolSelect = `SELECT id::text,code,name,description,endpoint_url,http_method,input_schema,timeout_seconds,department_ids::text[],rate_limit_rpm,approval_required,enabled,octet_length(encrypted_headers)>0,revision,created_at,updated_at,encrypted_headers,headers_kek_version FROM gateway.tool_definitions`
|
|
|
|
func scanTool(row pgx.Row) (Tool, error) {
|
|
var t Tool
|
|
err := row.Scan(&t.ID, &t.Code, &t.Name, &t.Description, &t.EndpointURL, &t.HTTPMethod, &t.InputSchema, &t.TimeoutSeconds, &t.DepartmentIDs, &t.RateLimitRPM, &t.ApprovalRequired, &t.Enabled, &t.HasSecretHeaders, &t.Revision, &t.CreatedAt, &t.UpdatedAt, &t.EncryptedHeaders, &t.HeadersKEKVersion)
|
|
return t, mapNotFound(err)
|
|
}
|
|
func (s *ToolService) List(ctx context.Context) ([]Tool, error) {
|
|
rows, err := s.assets.pool.Query(ctx, toolSelect+` ORDER BY updated_at DESC`)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
items := []Tool{}
|
|
for rows.Next() {
|
|
t, err := scanTool(rows)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, t)
|
|
}
|
|
return items, rows.Err()
|
|
}
|
|
func (s *ToolService) Get(ctx context.Context, id string) (Tool, error) {
|
|
return scanTool(s.assets.pool.QueryRow(ctx, toolSelect+` WHERE id=$1`, id))
|
|
}
|
|
func (s *ToolService) GetByCode(ctx context.Context, code string) (Tool, error) {
|
|
return scanTool(s.assets.pool.QueryRow(ctx, toolSelect+` WHERE code=$1 AND enabled`, code))
|
|
}
|
|
|
|
func (s *ToolService) Save(ctx context.Context, id string, input ToolInput, actorID string, create bool) (Tool, error) {
|
|
if err := s.validate(ctx, &input, create); err != nil {
|
|
return Tool{}, err
|
|
}
|
|
tx, err := s.assets.pool.Begin(ctx)
|
|
if err != nil {
|
|
return Tool{}, err
|
|
}
|
|
defer rollback(ctx, tx)
|
|
var encrypted []byte
|
|
var version int
|
|
if input.Headers != nil {
|
|
raw, _ := json.Marshal(input.Headers)
|
|
encrypted, version, err = s.cipher.Encrypt(raw)
|
|
if err != nil {
|
|
return Tool{}, err
|
|
}
|
|
}
|
|
if create {
|
|
id, err = newUUID()
|
|
if err != nil {
|
|
return Tool{}, err
|
|
}
|
|
_, err = tx.Exec(ctx, `INSERT INTO gateway.tool_definitions(id,code,name,description,endpoint_url,http_method,encrypted_headers,headers_kek_version,input_schema,timeout_seconds,department_ids,rate_limit_rpm,approval_required,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15)`, id, input.Code, input.Name, input.Description, input.EndpointURL, input.HTTPMethod, encrypted, version, input.InputSchema, input.TimeoutSeconds, input.DepartmentIDs, input.RateLimitRPM, input.ApprovalRequired, input.Enabled, actorID)
|
|
} else {
|
|
if input.Headers == nil {
|
|
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.tool_definitions SET code=$2,name=$3,description=$4,endpoint_url=$5,http_method=$6,input_schema=$7,timeout_seconds=$8,department_ids=$9,rate_limit_rpm=$10,approval_required=$11,enabled=$12,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.EndpointURL, input.HTTPMethod, input.InputSchema, input.TimeoutSeconds, input.DepartmentIDs, input.RateLimitRPM, input.ApprovalRequired, input.Enabled)
|
|
err = updateErr
|
|
if err == nil && tag.RowsAffected() == 0 {
|
|
return Tool{}, ErrNotFound
|
|
}
|
|
} else {
|
|
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.tool_definitions SET code=$2,name=$3,description=$4,endpoint_url=$5,http_method=$6,encrypted_headers=$7,headers_kek_version=$8,input_schema=$9,timeout_seconds=$10,department_ids=$11,rate_limit_rpm=$12,approval_required=$13,enabled=$14,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.EndpointURL, input.HTTPMethod, encrypted, version, input.InputSchema, input.TimeoutSeconds, input.DepartmentIDs, input.RateLimitRPM, input.ApprovalRequired, input.Enabled)
|
|
err = updateErr
|
|
if err == nil && tag.RowsAffected() == 0 {
|
|
return Tool{}, ErrNotFound
|
|
}
|
|
}
|
|
}
|
|
if err != nil {
|
|
return Tool{}, err
|
|
}
|
|
event := "tool.updated"
|
|
if create {
|
|
event = "tool.created"
|
|
}
|
|
if err = emit(ctx, tx, event, "tool", id, actorID, nil); err != nil {
|
|
return Tool{}, err
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return Tool{}, err
|
|
}
|
|
return s.Get(ctx, id)
|
|
}
|
|
|
|
func (s *ToolService) Delete(ctx context.Context, id, actorID string) error {
|
|
tx, err := s.assets.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer rollback(ctx, tx)
|
|
var used bool
|
|
if err = tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.applications WHERE draft_config->'tool_ids' ? $1 UNION ALL SELECT 1 FROM gateway.application_versions WHERE config->'tool_ids' ? $1)`, id).Scan(&used); err != nil {
|
|
return err
|
|
}
|
|
if used {
|
|
return ErrConflict
|
|
}
|
|
tag, err := tx.Exec(ctx, `DELETE FROM gateway.tool_definitions WHERE id=$1`, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if tag.RowsAffected() == 0 {
|
|
return ErrNotFound
|
|
}
|
|
if err = emit(ctx, tx, "tool.deleted", "tool", id, actorID, nil); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (s *ToolService) headers(tool Tool) (map[string]string, error) {
|
|
plain, err := s.cipher.Decrypt(tool.EncryptedHeaders, tool.HeadersKEKVersion)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
headers := map[string]string{}
|
|
if err = json.Unmarshal(plain, &headers); err != nil {
|
|
return nil, errors.New("工具请求头密文内容无效")
|
|
}
|
|
return headers, nil
|
|
}
|
|
|
|
// ErrToolApprovalRequired 表示工具需管理员审批后才能调用。
|
|
var ErrToolApprovalRequired = errors.New("工具需要管理员审批后才能调用")
|
|
|
|
// ErrToolRateLimited 表示工具调用频率超限。
|
|
var ErrToolRateLimited = errors.New("工具调用频率超限，请稍后重试")
|
|
|
|
// enforceGovernance 在工具执行前做治理校验:审批标记 + 调用频率上限。
|
|
// 审批缺失时自动发起一次申请(每工具至多一个待审项);限流用固定窗口原子
|
|
// upsert,多实例共享同一额度。apiKeyID 对应的门户用户若配置了个人智能体
|
|
// 安全策略:auto_approve_tools 跳过审批门,rate_limit_multiplier 按倍数
|
|
// 放宽个人限流(个人窗口独立计数)。返回 (allowed, err)。
|
|
func (s *ToolService) enforceGovernance(ctx context.Context, tool Tool, apiKeyID string) (bool, error) {
|
|
if s == nil || s.assets == nil || s.assets.pool == nil {
|
|
return false, errors.New("工具服务不可用")
|
|
}
|
|
// 解析调用者的个人策略(仅门户运行时凭据可能命中)。
|
|
portalUserID, policy, err := s.personalPolicy(ctx, apiKeyID)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if tool.ApprovalRequired && !policy.AutoApproveTools {
|
|
var approved bool
|
|
if err := s.assets.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.tool_approval_requests WHERE tool_id=$1 AND status='approved')`, tool.ID).Scan(&approved); err != nil {
|
|
return false, err
|
|
}
|
|
if !approved {
|
|
// 自动发起待审申请(唯一部分索引防重复),通知管理员。
|
|
requestID, err := newUUID()
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
eventID, err := newUUID()
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
tx, err := s.assets.pool.Begin(ctx)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
tag, err := tx.Exec(ctx, `INSERT INTO gateway.tool_approval_requests(id,tool_id,reason) VALUES($1,$2,$3) ON CONFLICT DO NOTHING`, requestID, tool.ID, "工具首次调用,自动发起审批")
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if tag.RowsAffected() > 0 {
|
|
payload, _ := json.Marshal(map[string]any{"tool_id": tool.ID, "tool_code": tool.Code, "request_id": requestID})
|
|
if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,'tool_approval.requested',1,'tool',$2,$3)`, eventID, tool.ID, payload); err != nil {
|
|
return false, err
|
|
}
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return false, err
|
|
}
|
|
return false, fmt.Errorf("%w: %s（已自动发起审批）", ErrToolApprovalRequired, tool.Name)
|
|
}
|
|
}
|
|
limit := tool.RateLimitRPM
|
|
if portalUserID != "" && policy.RateLimitMultiplier > 1 {
|
|
limit = tool.RateLimitRPM * policy.RateLimitMultiplier
|
|
}
|
|
if limit > 0 {
|
|
userKey := portalUserID
|
|
if userKey == "" {
|
|
userKey = personalPolicySentinel
|
|
}
|
|
var count int64
|
|
err := s.assets.pool.QueryRow(ctx, `INSERT INTO gateway.tool_rate_usage(tool_id,portal_user_id,window_start,call_count)
|
|
VALUES($1,$2,date_trunc('minute',clock_timestamp()),1)
|
|
ON CONFLICT (tool_id,portal_user_id,window_start) DO UPDATE SET call_count=gateway.tool_rate_usage.call_count+1
|
|
RETURNING call_count`, tool.ID, userKey).Scan(&count)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if count > int64(limit) {
|
|
return false, ErrToolRateLimited
|
|
}
|
|
}
|
|
return true, nil
|
|
}
|
|
|
|
// personalPolicySentinel 是非个人调用的限流窗口占位键(全零 UUID)。
|
|
const personalPolicySentinel = "00000000-0000-0000-0000-000000000000"
|
|
|
|
// personalPolicy 解析 API Key 对应的门户用户及其个人智能体安全策略。
|
|
// 非门户 Key 返回空 userID 与默认策略(不跳过审批、倍数 1)。
|
|
func (s *ToolService) personalPolicy(ctx context.Context, apiKeyID string) (string, AgentPolicy, error) {
|
|
if strings.TrimSpace(apiKeyID) == "" {
|
|
return "", AgentPolicy{RateLimitMultiplier: 1}, nil
|
|
}
|
|
var portalUserID *string
|
|
err := s.assets.pool.QueryRow(ctx, `SELECT portal_user_id::text FROM gateway.api_keys WHERE id=$1`, apiKeyID).Scan(&portalUserID)
|
|
if err != nil || portalUserID == nil || *portalUserID == "" {
|
|
return "", AgentPolicy{RateLimitMultiplier: 1}, nil
|
|
}
|
|
var policy AgentPolicy
|
|
err = s.assets.pool.QueryRow(ctx, `SELECT auto_approve_tools,rate_limit_multiplier FROM gateway.portal_agent_policies WHERE portal_user_id=$1`, *portalUserID).Scan(&policy.AutoApproveTools, &policy.RateLimitMultiplier)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
policy.RateLimitMultiplier = 1
|
|
return *portalUserID, policy, nil
|
|
}
|
|
if err != nil {
|
|
return "", AgentPolicy{}, err
|
|
}
|
|
return *portalUserID, policy, nil
|
|
}
|
|
|
|
// ListApprovalRequests 返回工具审批申请(含工具信息)。
|
|
func (s *ToolService) ListApprovalRequests(ctx context.Context, status string) ([]map[string]any, error) {
|
|
if s == nil || s.assets == nil || s.assets.pool == nil {
|
|
return nil, errors.New("工具服务不可用")
|
|
}
|
|
where, args := " WHERE true", []any{}
|
|
if status != "" {
|
|
args = append(args, status)
|
|
where += fmt.Sprintf(" AND r.status=$%d", len(args))
|
|
}
|
|
rows, err := s.assets.pool.Query(ctx, `SELECT r.id::text,t.code,t.name,r.status,r.reason,r.decision_note,r.created_at,r.decided_at,r.decided_by::text FROM gateway.tool_approval_requests r JOIN gateway.tool_definitions t ON t.id=r.tool_id`+where+` ORDER BY r.created_at DESC LIMIT 200`, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
items := []map[string]any{}
|
|
for rows.Next() {
|
|
var id, code, name, status, reason, note string
|
|
var decidedAt, decidedBy any
|
|
var createdAt any
|
|
if err := rows.Scan(&id, &code, &name, &status, &reason, ¬e, &createdAt, &decidedAt, &decidedBy); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, map[string]any{"id": id, "tool_code": code, "tool_name": name, "status": status, "reason": reason, "decision_note": note, "created_at": createdAt, "decided_at": decidedAt, "decided_by": decidedBy})
|
|
}
|
|
return items, rows.Err()
|
|
}
|
|
|
|
// DecideApprovalRequest 审批工具申请;通过后工具立即可调用。
|
|
func (s *ToolService) DecideApprovalRequest(ctx context.Context, id, status, note, actorID string) error {
|
|
if status != "approved" && status != "rejected" {
|
|
return errors.New("审批状态无效")
|
|
}
|
|
if len(note) > 4000 {
|
|
return errors.New("审批备注过长")
|
|
}
|
|
tx, err := s.assets.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
var toolID string
|
|
err = tx.QueryRow(ctx, `UPDATE gateway.tool_approval_requests SET status=$2,decision_note=$3,decided_by=$4,decided_at=clock_timestamp() WHERE id=$1 AND status='pending' RETURNING tool_id::text`, id, status, note, actorID).Scan(&toolID)
|
|
if err != nil {
|
|
return mapNotFound(err)
|
|
}
|
|
eventID, _ := newUUID()
|
|
payload, _ := json.Marshal(map[string]any{"request_id": id, "tool_id": toolID, "status": status, "actor_id": actorID})
|
|
if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,'tool_approval.decided',1,'tool',$2,$3)`, eventID, toolID, payload); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (s *ToolService) Execute(ctx context.Context, tool Tool, input map[string]any, apiKeyID, requestID string) (result map[string]any, err error) {
|
|
started := time.Now()
|
|
status := "success"
|
|
var responseStatus *int
|
|
defer func() {
|
|
if err != nil {
|
|
status = "error"
|
|
}
|
|
runID, idErr := newUUID()
|
|
if idErr == nil {
|
|
message := ""
|
|
if err != nil {
|
|
message = err.Error()
|
|
if len(message) > 1000 {
|
|
message = message[:1000]
|
|
}
|
|
}
|
|
_, _ = s.assets.pool.Exec(context.WithoutCancel(ctx), `INSERT INTO gateway.tool_runs(id,tool_id,api_key_id,request_id,status,response_status,latency_ms,error) VALUES($1,$2,nullif($3,'')::uuid,$4,$5,$6,$7,$8)`, runID, tool.ID, apiKeyID, requestID, status, responseStatus, time.Since(started).Milliseconds(), message)
|
|
}
|
|
}()
|
|
if err = validateToolInput(tool.InputSchema, input); err != nil {
|
|
return nil, err
|
|
}
|
|
// 治理校验:审批标记 + 频率上限(个人策略可跳过审批/放宽限流)。被拒时
|
|
// 不算一次成功调用(但会记 tool_runs 失败)。
|
|
if allowed, governanceErr := s.enforceGovernance(ctx, tool, apiKeyID); !allowed {
|
|
return nil, governanceErr
|
|
}
|
|
headers, err := s.headers(tool)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
payload, err := json.Marshal(input)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var body io.Reader
|
|
parsed, err := url.Parse(tool.EndpointURL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if tool.HTTPMethod == http.MethodGet {
|
|
query := parsed.Query()
|
|
for key, value := range input {
|
|
query.Set(key, toString(value))
|
|
}
|
|
parsed.RawQuery = query.Encode()
|
|
} else {
|
|
body = bytes.NewReader(payload)
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, tool.HTTPMethod, parsed.String(), body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for key, value := range headers {
|
|
request.Header.Set(key, value)
|
|
}
|
|
request.Header.Set("Accept", "application/json")
|
|
if body != nil {
|
|
request.Header.Set("Content-Type", "application/json")
|
|
}
|
|
client := &http.Client{Timeout: time.Duration(tool.TimeoutSeconds) * time.Second, Transport: &http.Transport{DialContext: safeToolDial(s.allowPrivate), ForceAttemptHTTP2: true, TLSHandshakeTimeout: 5 * time.Second, ResponseHeaderTimeout: time.Duration(tool.TimeoutSeconds) * time.Second}, CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("工具端点不允许重定向") }}
|
|
response, err := client.Do(request)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("工具调用失败: %w", err)
|
|
}
|
|
defer response.Body.Close()
|
|
code := response.StatusCode
|
|
responseStatus = &code
|
|
raw, err := io.ReadAll(io.LimitReader(response.Body, (1<<20)+1))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(raw) > 1<<20 {
|
|
return nil, errors.New("工具响应超过 1 MiB")
|
|
}
|
|
var decoded any
|
|
if json.Unmarshal(raw, &decoded) != nil {
|
|
decoded = string(raw)
|
|
}
|
|
return map[string]any{"status_code": response.StatusCode, "body": decoded}, nil
|
|
}
|
|
|
|
func validateToolInput(raw json.RawMessage, input map[string]any) error {
|
|
var schema struct {
|
|
Required []string `json:"required"`
|
|
Properties map[string]struct {
|
|
Type string `json:"type"`
|
|
} `json:"properties"`
|
|
}
|
|
if len(raw) == 0 {
|
|
return nil
|
|
}
|
|
if err := json.Unmarshal(raw, &schema); err != nil {
|
|
return errors.New("工具 input_schema 无效")
|
|
}
|
|
for _, name := range schema.Required {
|
|
if _, ok := input[name]; !ok {
|
|
return fmt.Errorf("缺少工具必填参数 %s", name)
|
|
}
|
|
}
|
|
for name, property := range schema.Properties {
|
|
value, ok := input[name]
|
|
if !ok || property.Type == "" {
|
|
continue
|
|
}
|
|
valid := false
|
|
switch property.Type {
|
|
case "string":
|
|
_, valid = value.(string)
|
|
case "number":
|
|
switch value.(type) {
|
|
case float64, float32, int, int64, json.Number:
|
|
valid = true
|
|
}
|
|
case "integer":
|
|
switch v := value.(type) {
|
|
case int, int64:
|
|
valid = true
|
|
case float64:
|
|
valid = v == float64(int64(v))
|
|
}
|
|
case "boolean":
|
|
_, valid = value.(bool)
|
|
case "object":
|
|
_, valid = value.(map[string]any)
|
|
case "array":
|
|
_, valid = value.([]any)
|
|
}
|
|
if !valid {
|
|
return fmt.Errorf("工具参数 %s 类型应为 %s", name, property.Type)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func safeToolDial(allowPrivate bool) func(context.Context, string, string) (net.Conn, error) {
|
|
dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
|
|
if allowPrivate {
|
|
return dialer.DialContext
|
|
}
|
|
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
|
host, port, err := net.SplitHostPort(address)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(addresses) == 0 {
|
|
return nil, errors.New("工具主机没有解析结果")
|
|
}
|
|
for _, candidate := range addresses {
|
|
ip := candidate.IP
|
|
if ip == nil || ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsMulticast() || ip.IsUnspecified() {
|
|
return nil, fmt.Errorf("工具主机解析到受限地址 %s", ip)
|
|
}
|
|
}
|
|
return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
|
|
}
|
|
}
|