ai-gateway-go/internal/workbench/runtime_marketplace.go

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"`
	SkillIDs     []string         `json:"skill_ids"`
	MCPServerIDs []string         `json:"mcp_server_ids"`
}

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
	modelCallCount := 0
	conversationID := strings.TrimSpace(r.Header.Get("X-Gateway-Conversation-ID"))
	traceID := h.beginTrace(r.Context(), principal, "digital_employee", employee.ID, employee.Code, conversationID)
	defer func() {
		traceCtx := context.WithoutCancel(r.Context())
		h.finishTrace(traceCtx, traceID, status, runError, retrievalCount, modelCallCount, toolCount)
		runID, idErr := newUUID()
		if idErr == nil {
			_, _ = h.service.pool.Exec(traceCtx, `INSERT INTO gateway.digital_employee_runs(id,digital_employee_id,api_key_id,trace_id,request_id,status,latency_ms,retrieval_count,tool_count,error) VALUES($1,$2,nullif($3,'')::uuid,nullif($4,'')::uuid,$5,$6,$7,$8,$9,$10)`, runID, employee.ID, principal.APIKeyID, traceID, gateway.RequestID(r.Context()), status, time.Since(started).Milliseconds(), retrievalCount, toolCount, runError)
		}
	}()
	executors, payload, prepareErr := h.prepareDigitalEmployee(r.Context(), employee, input, principal, &retrievalCount, portalUserID, traceID)
	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++ {
		modelCallCount++
		statusCode, responseHeaders, response, err = h.callGatewayWithTrace(r, payload, traceID, round)
		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 := h.executeToolWithTrace(r.Context(), traceID, call.Name, call.ID, round, func() (map[string]any, error) {
				return 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++
		}
	}
	// 事实核验与普通应用一致:block 策略不能因走数字员工入口而被绕过。
	if h.factCheck != nil {
		h.applyFactCheck(r, applicationRequest{Messages: input.Messages, Variables: input.Variables}, response, employee.DepartmentIDs)
	}
	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, traceID 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)
	}
	selectedSkills, err := selectedBindings(input.SkillIDs, employee.SkillIDs, "Skill")
	if err != nil {
		return nil, nil, err
	}
	selectedMCPServers, err := selectedBindings(input.MCPServerIDs, employee.MCPServerIDs, "MCP")
	if err != nil {
		return nil, nil, err
	}
	skills := map[string]Skill{}
	for _, skillID := range selectedSkills {
		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.searchWithTrace(ctx, traceID, 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 selectedMCPServers {
		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
}

// selectedBindings lets scheduled tasks restrict a digital employee to a
// subset of its published Skill/MCP bindings. Empty means the employee's full
// published binding set; callers can never add resources it does not own.
func selectedBindings(selected, allowed []string, label string) ([]string, error) {
	if len(selected) == 0 {
		return allowed, nil
	}
	allowedSet := make(map[string]bool, len(allowed))
	for _, id := range allowed {
		allowedSet[id] = true
	}
	seen := map[string]bool{}
	result := make([]string, 0, len(selected))
	for _, id := range selected {
		if !allowedSet[id] {
			return nil, fmt.Errorf("请求的 %s %s 未绑定到数字员工", label, id)
		}
		if !seen[id] {
			seen[id] = true
			result = append(result, id)
		}
	}
	return result, 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)
}