ai-gateway-go/internal/workbench/tools.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 秒之间")
	}
	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[],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.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,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13)`, id, input.Code, input.Name, input.Description, input.EndpointURL, input.HTTPMethod, encrypted, version, input.InputSchema, input.TimeoutSeconds, input.DepartmentIDs, 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,enabled=$10,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.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,enabled=$12,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.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
}

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
	}
	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))
	}
}