Files
ai-gateway-go/internal/workbench/runtime_marketplace.go
T
superidou 5759c1862e AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。
含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 11:45:54 +08:00

439 lines
15 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package workbench
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"time"
"aigateway.local/core/internal/apikey"
"aigateway.local/core/internal/gateway"
)
// toolExecutor runs one tool (regular HTTP tool or MCP tool) during a digital
// employee chat round and returns the tool message body.
type toolExecutor func(ctx context.Context, args map[string]any) (map[string]any, error)
func (h *RuntimeHTTPHandler) renderSkill(w http.ResponseWriter, r *http.Request) {
principal, ok := h.principal(w, r)
if !ok {
return
}
var input struct {
Variables map[string]any `json:"variables"`
}
if !decodeRuntime(w, r, &input) {
return
}
skill, err := h.market.Skills.GetPublishedByCode(r.Context(), r.PathValue("code"))
if err != nil || !visible(skill.DepartmentIDs, principal, false) {
runtimeError(w, 404, "Skill 不存在或不可见")
return
}
rendered, err := h.market.Skills.Render(skill, input.Variables)
if err != nil {
runtimeError(w, 400, err.Error())
return
}
writeRuntime(w, 200, map[string]any{"code": skill.Code, "name": skill.Name, "rendered": rendered})
}
func (h *RuntimeHTTPHandler) listMCPServers(w http.ResponseWriter, r *http.Request) {
principal, ok := h.principal(w, r)
if !ok {
return
}
portalUserID, err := h.portalUserID(r.Context(), principal)
if err != nil {
runtimeError(w, 503, "安装信息暂不可用")
return
}
servers, err := h.market.MCPServers.List(r.Context())
if err != nil {
runtimeError(w, 503, "MCP 服务器服务暂不可用")
return
}
result := []map[string]any{}
for _, server := range servers {
if server.Status != "published" || !server.Enabled {
continue
}
allowed, aerr := h.mcpAccessAllowed(r.Context(), principal, portalUserID, server)
if aerr != nil || !allowed {
continue
}
result = append(result, map[string]any{"code": server.Code, "name": server.Name, "description": server.Description, "transport": server.Transport})
}
writeRuntime(w, 200, map[string]any{"object": "list", "data": result})
}
func (h *RuntimeHTTPHandler) mcpServerTools(w http.ResponseWriter, r *http.Request) {
principal, ok := h.principal(w, r)
if !ok {
return
}
server, err := h.market.MCPServers.GetPublishedByCode(r.Context(), r.PathValue("code"))
if err != nil || !h.mcpAccessible(r, principal, server) {
runtimeError(w, 404, "MCP 服务器不存在或不可访问")
return
}
headers, err := h.market.MCPServers.Headers(server)
if err != nil {
runtimeError(w, 503, "MCP 服务器凭据不可用")
return
}
tools, err := h.market.MCPClient.DiscoverTools(r.Context(), server, headers)
if err != nil {
runtimeError(w, 502, err.Error())
return
}
result := []map[string]any{}
for _, tool := range tools {
result = append(result, map[string]any{"name": mcpToolName(server.Code, tool.Name), "description": tool.Description})
}
writeRuntime(w, 200, map[string]any{"code": server.Code, "server": server.Name, "data": result})
}
func (h *RuntimeHTTPHandler) invokeMCPTool(w http.ResponseWriter, r *http.Request) {
principal, ok := h.principal(w, r)
if !ok {
return
}
server, err := h.market.MCPServers.GetPublishedByCode(r.Context(), r.PathValue("code"))
if err != nil || !h.mcpAccessible(r, principal, server) {
runtimeError(w, 404, "MCP 服务器不存在或不可访问")
return
}
var input struct {
Input map[string]any `json:"input"`
}
if !decodeRuntime(w, r, &input) {
return
}
headers, err := h.market.MCPServers.Headers(server)
if err != nil {
runtimeError(w, 503, "MCP 服务器凭据不可用")
return
}
result, err := h.market.MCPClient.CallTool(r.Context(), server, headers, r.PathValue("tool"), input.Input)
if err != nil {
runtimeError(w, 502, err.Error())
return
}
status := http.StatusOK
if result.IsError {
status = http.StatusBadGateway
}
writeRuntime(w, status, map[string]any{"name": mcpToolName(server.Code, r.PathValue("tool")), "content": result.Content, "is_error": result.IsError})
}
type digitalEmployeeRequest struct {
Messages []map[string]any `json:"messages"`
Variables map[string]any `json:"variables"`
}
func (h *RuntimeHTTPHandler) runDigitalEmployee(w http.ResponseWriter, r *http.Request) {
principal, ok := h.principal(w, r)
if !ok {
return
}
var input digitalEmployeeRequest
if !decodeRuntime(w, r, &input) {
return
}
employee, err := h.market.Employees.GetPublishedByCode(r.Context(), r.PathValue("code"))
if err != nil {
runtimeError(w, 404, "数字员工不存在或未发布")
return
}
portalUserID, err := h.portalUserID(r.Context(), principal)
if err != nil {
runtimeError(w, 503, "安装信息暂不可用")
return
}
// Department-scoped digital employees are restricted to members of the
// department; everyone else must have installed the resource first.
if !visible(employee.DepartmentIDs, principal, false) {
installed, ierr := h.market.Market.Installed(r.Context(), "digital_employee", employee.ID, portalUserID)
if ierr != nil || !installed {
runtimeError(w, 403, "未安装此数字员工")
return
}
}
started := time.Now()
status := "error"
runError := ""
retrievalCount := 0
toolCount := 0
defer func() {
runID, idErr := newUUID()
if idErr == nil {
_, _ = h.service.pool.Exec(context.WithoutCancel(r.Context()), `INSERT INTO gateway.digital_employee_runs(id,digital_employee_id,api_key_id,request_id,status,latency_ms,retrieval_count,tool_count,error) VALUES($1,$2,nullif($3,'')::uuid,$4,$5,$6,$7,$8,$9)`, runID, employee.ID, principal.APIKeyID, gateway.RequestID(r.Context()), status, time.Since(started).Milliseconds(), retrievalCount, toolCount, runError)
}
}()
executors, payload, prepareErr := h.prepareDigitalEmployee(r.Context(), employee, input, principal, &retrievalCount, portalUserID)
if prepareErr != nil {
runError = prepareErr.Error()
runtimeError(w, 400, runError)
return
}
var response map[string]any
var responseHeaders http.Header
var statusCode int
for round := 0; ; round++ {
statusCode, responseHeaders, response, err = h.callGateway(r, payload)
if err != nil {
runError = err.Error()
copyHeaders(w.Header(), responseHeaders)
runtimeError(w, statusCode, runError)
return
}
calls := extractToolCalls(response)
if len(calls) == 0 {
break
}
if round >= employee.MaxToolRounds {
runError = "工具调用轮次已达上限"
runtimeError(w, 502, runError)
return
}
choice := firstChoiceMessage(response)
payload["messages"] = append(payload["messages"].([]map[string]any), choice)
for _, call := range calls {
exec, exists := executors[call.Name]
if !exists {
runError = "模型请求了未授权工具 " + call.Name
runtimeError(w, 400, runError)
return
}
var args map[string]any
if json.Unmarshal([]byte(call.Arguments), &args) != nil {
args = map[string]any{}
}
result, executeErr := exec(r.Context(), args)
if executeErr != nil {
runError = executeErr.Error()
runtimeError(w, 502, runError)
return
}
encoded, _ := json.Marshal(result["body"])
payload["messages"] = append(payload["messages"].([]map[string]any), map[string]any{"role": "tool", "tool_call_id": call.ID, "name": call.Name, "content": string(encoded)})
toolCount++
}
}
response["digital_employee"] = map[string]any{"code": employee.Code, "name": employee.Name, "persona": employee.Persona, "retrieval_count": retrievalCount, "tool_calls": toolCount}
status = "success"
copyHeaders(w.Header(), responseHeaders)
writeRuntime(w, statusCode, response)
}
// prepareDigitalEmployee assembles the chat payload for a digital employee:
// persona + rendered skills as system context, knowledge RAG evidence, and the
// union of bound tools (regular + MCP) exposed to the model. It returns the
// tool executors keyed by the exact schema name the model may call.
func (h *RuntimeHTTPHandler) prepareDigitalEmployee(ctx context.Context, employee DigitalEmployee, input digitalEmployeeRequest, principal apikey.Principal, retrievalCount *int, portalUserID string) (map[string]toolExecutor, map[string]any, error) {
messages := make([]map[string]any, 0, len(input.Messages)+3)
total := 0
lastQuestion := ""
for _, message := range input.Messages {
role, _ := message["role"].(string)
content, contentOK := message["content"].(string)
if (role != "user" && role != "assistant") || !contentOK {
return nil, nil, errors.New("数字员工对话只接受 user/assistant 文本消息")
}
total += len(content)
if total > 100000 {
return nil, nil, errors.New("对话历史超过 100000 字符")
}
messages = append(messages, map[string]any{"role": role, "content": content})
if role == "user" {
lastQuestion = content
}
}
if lastQuestion == "" {
return nil, nil, errors.New("至少需要一条用户消息")
}
system := []string{}
if strings.TrimSpace(employee.Persona) != "" {
system = append(system, employee.Persona)
}
skills := map[string]Skill{}
for _, skillID := range employee.SkillIDs {
skill, err := h.market.Skills.Get(ctx, skillID)
if err != nil || !skill.Enabled || !visible(skill.DepartmentIDs, principal, false) {
return nil, nil, fmt.Errorf("绑定的 Skill %s 当前不可用", skillID)
}
rendered, err := h.market.Skills.Render(skill, input.Variables)
if err != nil {
return nil, nil, err
}
system = append(system, rendered)
skills[skillID] = skill
}
// Knowledge RAG across the employee's own bases and the bases bound to its
// skills (deduplicated).
evidence := []string{}
kbSeen := map[string]bool{}
rag := func(kbID string) error {
if kbSeen[kbID] {
return nil
}
kbSeen[kbID] = true
kb, err := h.service.GetKnowledgeBase(ctx, kbID)
if err != nil || !kb.Enabled || !visible(kb.DepartmentIDs, principal, false) {
return fmt.Errorf("绑定的知识库 %s 当前不可用", kbID)
}
hits, searchErr := h.retriever.Search(ctx, kbID, lastQuestion, employee.RetrievalTopK)
if searchErr != nil {
return nil
}
for _, hit := range hits {
*retrievalCount++
evidence = append(evidence, fmt.Sprintf("[资料%d%s]\n%s", len(evidence)+1, hit.DocumentTitle, hit.Content))
}
return nil
}
for _, kbID := range employee.KnowledgeBaseIDs {
if err := rag(kbID); err != nil {
return nil, nil, err
}
}
for _, skill := range skills {
for _, kbID := range skill.KnowledgeBaseIDs {
if err := rag(kbID); err != nil {
return nil, nil, err
}
}
}
if len(evidence) > 0 {
system = append(system, "请优先依据以下企业资料回答;资料不足时明确说明不确定,不得编造。引用时使用[资料N]。\n\n"+strings.Join(evidence, "\n\n"))
}
executors := map[string]toolExecutor{}
schemas := []map[string]any{}
addTool := func(toolID string) error {
tool, err := h.tools.Get(ctx, toolID)
if err != nil || !tool.Enabled || !visible(tool.DepartmentIDs, principal, true) {
return fmt.Errorf("绑定的工具 %s 当前不可用", toolID)
}
if _, exists := executors[tool.Code]; exists {
return nil
}
executors[tool.Code] = func(ctx context.Context, args map[string]any) (map[string]any, error) {
return h.tools.Execute(ctx, tool, args, principal.APIKeyID, gateway.RequestID(ctx))
}
var schema any
_ = json.Unmarshal(tool.InputSchema, &schema)
schemas = append(schemas, map[string]any{"type": "function", "function": map[string]any{"name": tool.Code, "description": tool.Description, "parameters": schema}})
return nil
}
addMCP := func(serverID string) error {
server, err := h.market.MCPServers.Get(ctx, serverID)
if err != nil || !server.Enabled {
return fmt.Errorf("绑定的 MCP 服务器 %s 当前不可用", serverID)
}
allowed, aerr := h.mcpAccessAllowed(ctx, principal, portalUserID, server)
if aerr != nil || !allowed {
return fmt.Errorf("绑定的 MCP 服务器 %s 不可访问", server.Code)
}
headers, err := h.market.MCPServers.Headers(server)
if err != nil {
return fmt.Errorf("绑定的 MCP 服务器 %s 凭据不可用", server.Code)
}
mcpTools, err := h.market.MCPClient.DiscoverTools(ctx, server, headers)
if err != nil {
// A bound server that is transiently unreachable must not brick the
// whole chat; skip its tools and let the employee degrade.
h.logger.Warn("digital employee MCP discovery failed", "server", server.Code, "error", err)
return nil
}
for _, tool := range mcpTools {
name := mcpToolName(server.Code, tool.Name)
if _, exists := executors[name]; exists {
continue
}
executors[name] = func(ctx context.Context, args map[string]any) (map[string]any, error) {
result, err := h.market.MCPClient.CallTool(ctx, server, headers, tool.Name, args)
if err != nil {
return nil, err
}
if result.IsError {
return map[string]any{"body": "[MCP 工具执行失败]\n" + result.Content}, nil
}
return map[string]any{"body": result.Content}, nil
}
var schema any
_ = json.Unmarshal(tool.InputSchema, &schema)
schemas = append(schemas, map[string]any{"type": "function", "function": map[string]any{"name": name, "description": tool.Description, "parameters": schema}})
}
return nil
}
for _, toolID := range employee.ToolIDs {
if err := addTool(toolID); err != nil {
return nil, nil, err
}
}
for _, skill := range skills {
for _, toolID := range skill.ToolIDs {
if err := addTool(toolID); err != nil {
return nil, nil, err
}
}
}
for _, serverID := range employee.MCPServerIDs {
if err := addMCP(serverID); err != nil {
return nil, nil, err
}
}
for _, skill := range skills {
for _, serverID := range skill.MCPServerIDs {
if err := addMCP(serverID); err != nil {
return nil, nil, err
}
}
}
if len(system) > 0 {
messages = append([]map[string]any{{"role": "system", "content": strings.Join(system, "\n\n")}}, messages...)
}
payload := map[string]any{"model": employee.Model, "messages": messages, "stream": false, "temperature": employee.Temperature}
if len(schemas) > 0 {
payload["tools"] = schemas
payload["tool_choice"] = "auto"
}
return executors, payload, nil
}
// portalUserID resolves the portal user behind an API key, if any. Resource
// marketplace installs are scoped to portal users.
func (h *RuntimeHTTPHandler) portalUserID(ctx context.Context, principal apikey.Principal) (string, error) {
if principal.APIKeyID == "" {
return "", nil
}
userID, _, err := h.market.Market.PortalUserForAPIKey(ctx, principal.APIKeyID)
return userID, err
}
// mcpAccessible reports whether a principal may reach a published MCP server:
// either the server is department-visible to them, or they have installed it
// from the marketplace.
func (h *RuntimeHTTPHandler) mcpAccessible(r *http.Request, principal apikey.Principal, server MCPServer) bool {
portalUserID, err := h.portalUserID(r.Context(), principal)
if err != nil {
return false
}
allowed, err := h.mcpAccessAllowed(r.Context(), principal, portalUserID, server)
return err == nil && allowed
}
func (h *RuntimeHTTPHandler) mcpAccessAllowed(ctx context.Context, principal apikey.Principal, portalUserID string, server MCPServer) (bool, error) {
if visible(server.DepartmentIDs, principal, true) {
return true, nil
}
if portalUserID == "" {
return false, nil
}
return h.market.Market.Installed(ctx, "mcp_server", server.ID, portalUserID)
}