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>
This commit is contained in:
@@ -0,0 +1,783 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
type AdminHTTPHandler struct {
|
||||
service *Service
|
||||
tools *ToolService
|
||||
notifications *NotificationService
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewAdminHTTPHandler(service *Service, tools *ToolService, notifications *NotificationService, identityService *identity.Service) *AdminHTTPHandler {
|
||||
h := &AdminHTTPHandler{service: service, tools: tools, notifications: notifications, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/prompt-categories", h.listPromptCategories)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/prompt-categories", h.createPromptCategory)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/prompt-categories/{id}", h.deletePromptCategory)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/prompts", h.listPrompts)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/prompts", h.createPrompt)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/prompts/{id}", h.getPrompt)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/prompts/{id}", h.updatePrompt)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/prompts/{id}", h.deletePrompt)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/prompts/{id}/versions", h.listPromptVersions)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/prompts/{id}/versions", h.addPromptVersion)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/prompts/{id}/versions/{version}/activate", h.activatePromptVersion)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/prompts/{id}/render", h.previewPrompt)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/knowledge-bases", h.listKnowledgeBases)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/knowledge-bases", h.createKnowledgeBase)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/knowledge-bases/{id}", h.updateKnowledgeBase)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/knowledge-bases/{id}", h.deleteKnowledgeBase)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/knowledge-bases/{id}/documents", h.listDocuments)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/knowledge-bases/{id}/documents", h.addDocument)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/knowledge-bases/{id}/documents/{document_id}", h.deleteDocument)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/knowledge-bases/{id}/documents/{document_id}/reprocess", h.reprocessDocument)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/knowledge-bases/{id}/search", h.searchKnowledge)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/tools", h.listTools)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/tools", h.createTool)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/tools/{id}", h.updateTool)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/tools/{id}", h.deleteTool)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/tools/{id}/test", h.testTool)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/applications", h.listApplications)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/applications", h.createApplication)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/applications/catalog", h.applicationCatalog)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/applications/asset-dependencies", h.applicationDependencies)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/applications/{id}", h.getApplication)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/applications/{id}", h.updateApplication)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/applications/{id}", h.deleteApplication)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/applications/{id}/publish", h.publishApplication)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/applications/{id}/rollback/{version}", h.rollbackApplication)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/applications/{id}/versions", h.listApplicationVersions)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/applications/{id}/runs", h.listApplicationRuns)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/notification-channels", h.listChannels)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/notification-channels", h.createChannel)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/notification-channels/{id}", h.updateChannel)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/notification-channels/{id}", h.deleteChannel)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/notification-deliveries", h.listDeliveries)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/notification-deliveries/{id}/retry", h.retryDelivery)
|
||||
return h
|
||||
}
|
||||
func (h *AdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
|
||||
func (h *AdminHTTPHandler) require(w http.ResponseWriter, r *http.Request, permission string) (identity.Account, bool) {
|
||||
account, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
|
||||
if err != nil {
|
||||
apiresponse.Error(w, 401, "登录状态无效")
|
||||
return account, false
|
||||
}
|
||||
if !identity.HasPermission(account, permission) {
|
||||
apiresponse.Error(w, 403, "缺少 AI 资产管理权限")
|
||||
return account, false
|
||||
}
|
||||
return account, true
|
||||
}
|
||||
func decodeAsset(w http.ResponseWriter, r *http.Request, target any) bool {
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 3<<20))
|
||||
decoder.DisallowUnknownFields()
|
||||
if err := decoder.Decode(target); err != nil {
|
||||
apiresponse.Error(w, 400, "请求格式无效")
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
func assetError(w http.ResponseWriter, err error) {
|
||||
switch {
|
||||
case errors.Is(err, ErrNotFound):
|
||||
apiresponse.Error(w, 404, "资源不存在")
|
||||
case errors.Is(err, ErrConflict):
|
||||
apiresponse.Error(w, 409, "资源仍被已发布或草稿应用引用")
|
||||
default:
|
||||
apiresponse.Error(w, 400, err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listPromptCategories(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPromptRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListPromptCategories(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) createPromptCategory(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionPromptManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.CreatePromptCategory(r.Context(), input.Name, input.Description, a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) deletePromptCategory(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionPromptManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.service.DeletePromptCategory(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
type promptPayload struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
CategoryID *string `json:"category_id"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Content string `json:"content"`
|
||||
Variables []Variable `json:"variables"`
|
||||
ChangeNote string `json:"change_note"`
|
||||
}
|
||||
|
||||
func promptInput(p promptPayload) PromptInput {
|
||||
return PromptInput{Name: p.Name, Description: p.Description, CategoryID: p.CategoryID, Tags: p.Tags, DepartmentIDs: p.DepartmentIDs, Enabled: p.Enabled, Content: p.Content, Variables: p.Variables, ChangeNote: p.ChangeNote}
|
||||
}
|
||||
func (h *AdminHTTPHandler) listPrompts(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPromptRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListPrompts(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) getPrompt(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPromptRead); !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.service.GetPrompt(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) createPrompt(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionPromptManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p promptPayload
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.CreatePrompt(r.Context(), promptInput(p), a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) updatePrompt(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionPromptManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p promptPayload
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.UpdatePrompt(r.Context(), r.PathValue("id"), promptInput(p), a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) deletePrompt(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionPromptManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.service.DeletePrompt(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
func (h *AdminHTTPHandler) listPromptVersions(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPromptRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListPromptVersions(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) addPromptVersion(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionPromptManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p struct {
|
||||
Content string `json:"content"`
|
||||
Variables []Variable `json:"variables"`
|
||||
ChangeNote string `json:"change_note"`
|
||||
Activate bool `json:"activate"`
|
||||
}
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.AddPromptVersion(r.Context(), r.PathValue("id"), p.Content, p.Variables, p.ChangeNote, a.ID, p.Activate)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) activatePromptVersion(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionPromptManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
version, err := strconv.Atoi(r.PathValue("version"))
|
||||
if err != nil {
|
||||
apiresponse.Error(w, 400, "版本号无效")
|
||||
return
|
||||
}
|
||||
if err = h.service.ActivatePromptVersion(r.Context(), r.PathValue("id"), version, a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"activated": true})
|
||||
}
|
||||
func (h *AdminHTTPHandler) previewPrompt(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPromptRead); !ok {
|
||||
return
|
||||
}
|
||||
var p struct {
|
||||
Variables map[string]any `json:"variables"`
|
||||
}
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
prompt, err := h.service.GetPrompt(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
if prompt.Current == nil {
|
||||
apiresponse.Error(w, 409, "Prompt 没有生效版本")
|
||||
return
|
||||
}
|
||||
rendered, err := RenderPrompt(prompt.Current.Content, prompt.Current.Variables, p.Variables)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"rendered": rendered, "version": prompt.Current.Version})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listKnowledgeBases(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionKnowledgeRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListKnowledgeBases(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) createKnowledgeBase(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionKnowledgeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var k KnowledgeBase
|
||||
if !decodeAsset(w, r, &k) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.SaveKnowledgeBase(r.Context(), k, a.ID, true)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) updateKnowledgeBase(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionKnowledgeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var k KnowledgeBase
|
||||
if !decodeAsset(w, r, &k) {
|
||||
return
|
||||
}
|
||||
k.ID = r.PathValue("id")
|
||||
item, err := h.service.SaveKnowledgeBase(r.Context(), k, a.ID, false)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) deleteKnowledgeBase(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionKnowledgeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.service.DeleteKnowledgeBase(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
func (h *AdminHTTPHandler) listDocuments(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionKnowledgeRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListKnowledgeDocuments(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) addDocument(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionKnowledgeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p struct {
|
||||
Title string `json:"title"`
|
||||
SourceType string `json:"source_type"`
|
||||
SourceURI string `json:"source_uri"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.AddKnowledgeDocument(r.Context(), r.PathValue("id"), p.Title, p.SourceType, p.SourceURI, p.Content, a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) deleteDocument(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionKnowledgeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.service.DeleteKnowledgeDocument(r.Context(), r.PathValue("id"), r.PathValue("document_id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
func (h *AdminHTTPHandler) reprocessDocument(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionKnowledgeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.service.ReprocessKnowledgeDocument(r.Context(), r.PathValue("id"), r.PathValue("document_id"), a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) searchKnowledge(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionKnowledgeRead); !ok {
|
||||
return
|
||||
}
|
||||
var p struct {
|
||||
Query string `json:"query"`
|
||||
TopK int `json:"top_k"`
|
||||
}
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
items, err := NewPostgreSQLRetriever(h.service).Search(r.Context(), r.PathValue("id"), p.Query, p.TopK)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
type toolPayload struct {
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
EndpointURL string `json:"endpoint_url"`
|
||||
HTTPMethod string `json:"http_method"`
|
||||
Headers *map[string]string `json:"headers"`
|
||||
InputSchema json.RawMessage `json:"input_schema"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func toolInput(p toolPayload, create bool) ToolInput {
|
||||
var headers map[string]string
|
||||
if p.Headers != nil {
|
||||
headers = *p.Headers
|
||||
} else if create {
|
||||
headers = map[string]string{}
|
||||
}
|
||||
return ToolInput{Code: p.Code, Name: p.Name, Description: p.Description, EndpointURL: p.EndpointURL, HTTPMethod: p.HTTPMethod, Headers: headers, InputSchema: p.InputSchema, TimeoutSeconds: p.TimeoutSeconds, DepartmentIDs: p.DepartmentIDs, Enabled: p.Enabled}
|
||||
}
|
||||
func (h *AdminHTTPHandler) listTools(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionToolRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.tools.List(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) createTool(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionToolManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p toolPayload
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.tools.Save(r.Context(), "", toolInput(p, true), a.ID, true)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) updateTool(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionToolManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p toolPayload
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.tools.Save(r.Context(), r.PathValue("id"), toolInput(p, false), a.ID, false)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) deleteTool(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionToolManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.tools.Delete(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
func (h *AdminHTTPHandler) testTool(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionToolManage); !ok {
|
||||
return
|
||||
}
|
||||
var p struct {
|
||||
Input map[string]any `json:"input"`
|
||||
}
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
tool, err := h.tools.Get(r.Context(), r.PathValue("id"))
|
||||
if err == nil {
|
||||
result, executeErr := h.tools.Execute(r.Context(), tool, p.Input, "", "admin-test")
|
||||
err = executeErr
|
||||
if err == nil {
|
||||
apiresponse.OK(w, result)
|
||||
return
|
||||
}
|
||||
}
|
||||
assetError(w, err)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listApplications(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionApplicationRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListApplications(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) getApplication(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionApplicationRead); !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.service.GetApplication(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
versions, err := h.service.ListApplicationVersions(r.Context(), item.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"application": item, "versions": versions})
|
||||
}
|
||||
func (h *AdminHTTPHandler) applicationCatalog(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionApplicationRead); !ok {
|
||||
return
|
||||
}
|
||||
prompts, err := h.service.ListPrompts(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
knowledge, err := h.service.ListKnowledgeBases(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
tools, err := h.tools.List(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"prompts": prompts, "knowledge_bases": knowledge, "tools": tools})
|
||||
}
|
||||
func (h *AdminHTTPHandler) applicationDependencies(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionApplicationRead); !ok {
|
||||
return
|
||||
}
|
||||
assetType := strings.TrimSpace(r.URL.Query().Get("asset_type"))
|
||||
assetID := strings.TrimSpace(r.URL.Query().Get("asset_id"))
|
||||
if assetID == "" || (assetType != "prompt" && assetType != "knowledge" && assetType != "tool") {
|
||||
apiresponse.Error(w, 400, "asset_type 或 asset_id 无效")
|
||||
return
|
||||
}
|
||||
items, err := h.service.ApplicationDependencies(r.Context(), assetType, assetID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) createApplication(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionApplicationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var app Application
|
||||
if !decodeAsset(w, r, &app) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.SaveApplication(r.Context(), app, a.ID, true)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) updateApplication(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionApplicationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var app Application
|
||||
if !decodeAsset(w, r, &app) {
|
||||
return
|
||||
}
|
||||
app.ID = r.PathValue("id")
|
||||
item, err := h.service.SaveApplication(r.Context(), app, a.ID, false)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) deleteApplication(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionApplicationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.service.DeleteApplication(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
func (h *AdminHTTPHandler) publishApplication(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionApplicationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p struct {
|
||||
ChangeNote string `json:"change_note"`
|
||||
}
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
app, err := h.service.PublishApplication(r.Context(), r.PathValue("id"), p.ChangeNote, a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, app)
|
||||
}
|
||||
func (h *AdminHTTPHandler) rollbackApplication(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionApplicationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
version, err := strconv.Atoi(r.PathValue("version"))
|
||||
if err != nil || version < 1 {
|
||||
apiresponse.Error(w, 400, "版本号无效")
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
ChangeNote string `json:"change_note"`
|
||||
}
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.service.RollbackApplication(r.Context(), r.PathValue("id"), version, input.ChangeNote, a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) listApplicationVersions(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionApplicationRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListApplicationVersions(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) listApplicationRuns(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionApplicationRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListApplicationRuns(r.Context(), r.PathValue("id"), 100)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
type channelPayload struct {
|
||||
Name string `json:"name"`
|
||||
WebhookURL string `json:"webhook_url"`
|
||||
SigningSecret *string `json:"signing_secret"`
|
||||
EventPatterns []string `json:"event_patterns"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listChannels(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionNotificationRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.notifications.ListChannels(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) createChannel(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionNotificationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p channelPayload
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.notifications.SaveChannel(r.Context(), "", NotificationInput{Name: p.Name, WebhookURL: p.WebhookURL, SigningSecret: p.SigningSecret, EventPatterns: p.EventPatterns, Enabled: p.Enabled}, a.ID, true)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) updateChannel(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionNotificationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var p channelPayload
|
||||
if !decodeAsset(w, r, &p) {
|
||||
return
|
||||
}
|
||||
item, err := h.notifications.SaveChannel(r.Context(), r.PathValue("id"), NotificationInput{Name: p.Name, WebhookURL: p.WebhookURL, SigningSecret: p.SigningSecret, EventPatterns: p.EventPatterns, Enabled: p.Enabled}, a.ID, false)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
func (h *AdminHTTPHandler) deleteChannel(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionNotificationManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.notifications.DeleteChannel(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
func (h *AdminHTTPHandler) listDeliveries(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionNotificationRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.notifications.ListDeliveries(r.Context(), 100)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
func (h *AdminHTTPHandler) retryDelivery(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionNotificationManage); !ok {
|
||||
return
|
||||
}
|
||||
if err := h.notifications.RetryDelivery(r.Context(), r.PathValue("id")); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"retried": true})
|
||||
}
|
||||
|
||||
var _ = strings.TrimSpace
|
||||
@@ -0,0 +1,435 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
// MarketplaceAdminHTTPHandler exposes the resource-marketplace administration
|
||||
// surface: shared categories, MCP servers (with connectivity test), skills,
|
||||
// digital employees, and the unified catalog.
|
||||
type MarketplaceAdminHTTPHandler struct {
|
||||
market *MarketplaceService
|
||||
mcpServers *MCPServerService
|
||||
skills *SkillService
|
||||
employees *DigitalEmployeeService
|
||||
mcpClient *MCPClient
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewMarketplaceAdminHTTPHandler(market *MarketplaceService, mcpServers *MCPServerService, skills *SkillService, employees *DigitalEmployeeService, mcpClient *MCPClient, identityService *identity.Service) *MarketplaceAdminHTTPHandler {
|
||||
h := &MarketplaceAdminHTTPHandler{market: market, mcpServers: mcpServers, skills: skills, employees: employees, mcpClient: mcpClient, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/marketplace-categories", h.listMarketplaceCategories)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/marketplace-categories", h.createMarketplaceCategory)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/marketplace-categories/{id}", h.deleteMarketplaceCategory)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/mcp-servers", h.listMCPServers)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/mcp-servers", h.createMCPServer)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/mcp-servers/{id}", h.getMCPServer)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/mcp-servers/{id}", h.updateMCPServer)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/mcp-servers/{id}", h.deleteMCPServer)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/mcp-servers/{id}/test", h.testMCPServer)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/skills", h.listSkills)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/skills", h.createSkill)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/skills/{id}", h.getSkill)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/skills/{id}", h.updateSkill)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/skills/{id}", h.deleteSkill)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/skills/{id}/render", h.renderSkillPreview)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/digital-employees", h.listDigitalEmployees)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/digital-employees", h.createDigitalEmployee)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/digital-employees/{id}", h.getDigitalEmployee)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/digital-employees/{id}", h.updateDigitalEmployee)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/digital-employees/{id}", h.deleteDigitalEmployee)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/digital-employees/{id}/publish", h.publishDigitalEmployee)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/digital-employees/{id}/archive", h.archiveDigitalEmployee)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/digital-employees/{id}/runs", h.listDigitalEmployeeRuns)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/marketplace/catalog", h.marketplaceCatalog)
|
||||
return h
|
||||
}
|
||||
func (h *MarketplaceAdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
h.mux.ServeHTTP(w, r)
|
||||
}
|
||||
func (h *MarketplaceAdminHTTPHandler) require(w http.ResponseWriter, r *http.Request, permission string) (identity.Account, bool) {
|
||||
account, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
|
||||
if err != nil {
|
||||
apiresponse.Error(w, 401, "登录状态无效")
|
||||
return account, false
|
||||
}
|
||||
if !identity.HasPermission(account, permission) {
|
||||
apiresponse.Error(w, 403, "缺少资源市场管理权限")
|
||||
return account, false
|
||||
}
|
||||
return account, true
|
||||
}
|
||||
|
||||
// --- 分类 ---
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) listMarketplaceCategories(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionMarketplaceRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.market.ListCategories(r.Context(), r.URL.Query().Get("type"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) createMarketplaceCategory(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionMarketplaceManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
ResourceType string `json:"resource_type"`
|
||||
SortOrder int `json:"sort_order"`
|
||||
}
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.market.CreateCategory(r.Context(), Category{Name: input.Name, Description: input.Description, ResourceType: input.ResourceType, SortOrder: input.SortOrder}, a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) deleteMarketplaceCategory(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionMarketplaceManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.market.DeleteCategory(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
// --- MCP 服务器 ---
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) listMCPServers(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionMCPServerRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.mcpServers.List(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) createMCPServer(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionMCPServerManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input MCPServerInput
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.mcpServers.Save(r.Context(), "", input, a.ID, true)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) getMCPServer(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionMCPServerRead); !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.mcpServers.Get(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) updateMCPServer(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionMCPServerManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input MCPServerInput
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.mcpServers.Save(r.Context(), r.PathValue("id"), input, a.ID, false)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) deleteMCPServer(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionMCPServerManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.mcpServers.Delete(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
// testMCPServer verifies connectivity and lists the tools an MCP server
|
||||
// advertises so an operator can confirm it before publishing.
|
||||
func (h *MarketplaceAdminHTTPHandler) testMCPServer(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionMCPServerManage); !ok {
|
||||
return
|
||||
}
|
||||
server, err := h.mcpServers.Get(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
headers, err := h.mcpServers.Headers(server)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
tools, err := h.mcpClient.DiscoverTools(r.Context(), server, headers)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, 502, "MCP 服务器连通性测试失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, tools)
|
||||
}
|
||||
|
||||
// --- Skills ---
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) listSkills(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionSkillRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.skills.List(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) createSkill(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionSkillManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input SkillInput
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.skills.Save(r.Context(), "", input, a.ID, true)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) getSkill(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionSkillRead); !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.skills.Get(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) updateSkill(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionSkillManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input SkillInput
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.skills.Save(r.Context(), r.PathValue("id"), input, a.ID, false)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) deleteSkill(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionSkillManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.skills.Delete(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) renderSkillPreview(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionSkillRead); !ok {
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
Variables map[string]any `json:"variables"`
|
||||
}
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
skill, err := h.skills.Get(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
rendered, err := h.skills.Render(skill, input.Variables)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]any{"code": skill.Code, "name": skill.Name, "rendered": rendered})
|
||||
}
|
||||
|
||||
// --- 数字员工 ---
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) listDigitalEmployees(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionDigitalEmployeeRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.employees.List(r.Context())
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) createDigitalEmployee(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionDigitalEmployeeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input DigitalEmployeeInput
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.employees.Save(r.Context(), "", input, a.ID, true)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) getDigitalEmployee(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionDigitalEmployeeRead); !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.employees.Get(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) updateDigitalEmployee(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionDigitalEmployeeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input DigitalEmployeeInput
|
||||
if !decodeAsset(w, r, &input) {
|
||||
return
|
||||
}
|
||||
item, err := h.employees.Save(r.Context(), r.PathValue("id"), input, a.ID, false)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) deleteDigitalEmployee(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionDigitalEmployeeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := h.employees.Delete(r.Context(), r.PathValue("id"), a.ID); err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) publishDigitalEmployee(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionDigitalEmployeeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.employees.Publish(r.Context(), r.PathValue("id"), a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) archiveDigitalEmployee(w http.ResponseWriter, r *http.Request) {
|
||||
a, ok := h.require(w, r, identity.PermissionDigitalEmployeeManage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
item, err := h.employees.Archive(r.Context(), r.PathValue("id"), a.ID)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) listDigitalEmployeeRuns(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionDigitalEmployeeRead); !ok {
|
||||
return
|
||||
}
|
||||
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
|
||||
items, err := h.employees.ListRuns(r.Context(), r.PathValue("id"), limit)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
// --- 统一目录 ---
|
||||
|
||||
func (h *MarketplaceAdminHTTPHandler) marketplaceCatalog(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionMarketplaceRead); !ok {
|
||||
return
|
||||
}
|
||||
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
|
||||
query := r.URL.Query()
|
||||
items, err := h.market.Catalog(r.Context(), query.Get("type"), query.Get("category_id"), query.Get("tag"), query.Get("q"), limit)
|
||||
if err != nil {
|
||||
assetError(w, err)
|
||||
return
|
||||
}
|
||||
if items == nil {
|
||||
items = []MarketItem{}
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
@@ -0,0 +1,380 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
func validateApplication(app *Application) error {
|
||||
app.Code = strings.ToLower(strings.TrimSpace(app.Code))
|
||||
app.Name = strings.TrimSpace(app.Name)
|
||||
app.Description = strings.TrimSpace(app.Description)
|
||||
app.Status = strings.TrimSpace(app.Status)
|
||||
if !codeRE.MatchString(app.Code) || len(app.Code) < 3 || app.Name == "" || len(app.Name) > 128 || len(app.Description) > 4000 {
|
||||
return errors.New("应用编码、名称或描述格式无效")
|
||||
}
|
||||
switch app.Status {
|
||||
case "draft", "active", "suspended", "retired":
|
||||
case "":
|
||||
app.Status = "draft"
|
||||
default:
|
||||
return errors.New("应用状态无效")
|
||||
}
|
||||
var err error
|
||||
app.DepartmentIDs, err = normalizeStrings(app.DepartmentIDs, 100)
|
||||
return err
|
||||
}
|
||||
|
||||
func normalizeApplicationConfig(config *ApplicationConfig, requireModel bool) error {
|
||||
config.Model = strings.TrimSpace(config.Model)
|
||||
config.PromptTemplateID = strings.TrimSpace(config.PromptTemplateID)
|
||||
if requireModel && config.Model == "" {
|
||||
return errors.New("发布配置必须指定模型")
|
||||
}
|
||||
if len(config.Model) > 255 {
|
||||
return errors.New("模型名称过长")
|
||||
}
|
||||
var err error
|
||||
if config.KnowledgeBaseIDs, err = normalizeStrings(config.KnowledgeBaseIDs, 20); err != nil {
|
||||
return err
|
||||
}
|
||||
if config.ToolIDs, err = normalizeStrings(config.ToolIDs, 20); err != nil {
|
||||
return err
|
||||
}
|
||||
if config.RetrievalTopK == 0 {
|
||||
config.RetrievalTopK = 4
|
||||
}
|
||||
if config.RetrievalTopK < 1 || config.RetrievalTopK > 20 {
|
||||
return errors.New("retrieval_top_k 应在 1-20 之间")
|
||||
}
|
||||
if config.Temperature < 0 || config.Temperature > 2 {
|
||||
return errors.New("temperature 应在 0-2 之间")
|
||||
}
|
||||
if config.MaxToolRounds < 0 || config.MaxToolRounds > 8 {
|
||||
return errors.New("max_tool_rounds 应在 0-8 之间")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) validateApplicationRefs(ctx context.Context, config ApplicationConfig) error {
|
||||
if config.PromptTemplateID != "" {
|
||||
var ok bool
|
||||
if err := s.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.prompt_templates WHERE id=$1 AND enabled AND current_version IS NOT NULL)`, config.PromptTemplateID).Scan(&ok); err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return errors.New("绑定的 Prompt 不存在、未启用或没有生效版本")
|
||||
}
|
||||
}
|
||||
for _, id := range config.KnowledgeBaseIDs {
|
||||
var ok bool
|
||||
if err := s.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.knowledge_bases WHERE id=$1 AND enabled)`, id).Scan(&ok); err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return fmt.Errorf("知识库 %s 不存在或未启用", id)
|
||||
}
|
||||
}
|
||||
for _, id := range config.ToolIDs {
|
||||
var ok bool
|
||||
if err := s.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.tool_definitions WHERE id=$1 AND enabled)`, id).Scan(&ok); err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return fmt.Errorf("工具 %s 不存在或未启用", id)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const applicationSelect = `SELECT a.id::text,a.code,a.name,a.description,a.department_ids::text[],a.status,a.draft_config,a.published_version,p.config,a.revision,a.created_at,a.updated_at FROM gateway.applications a LEFT JOIN gateway.application_versions p ON p.application_id=a.id AND p.version=a.published_version`
|
||||
|
||||
func scanApplication(row pgx.Row) (Application, error) {
|
||||
var app Application
|
||||
var draft []byte
|
||||
var published []byte
|
||||
err := row.Scan(&app.ID, &app.Code, &app.Name, &app.Description, &app.DepartmentIDs, &app.Status, &draft, &app.PublishedVersion, &published, &app.Revision, &app.CreatedAt, &app.UpdatedAt)
|
||||
if err != nil {
|
||||
return app, mapNotFound(err)
|
||||
}
|
||||
_ = json.Unmarshal(draft, &app.DraftConfig)
|
||||
if len(published) > 0 {
|
||||
var config ApplicationConfig
|
||||
if json.Unmarshal(published, &config) == nil {
|
||||
app.PublishedConfig = &config
|
||||
}
|
||||
}
|
||||
return app, nil
|
||||
}
|
||||
func (s *Service) ListApplications(ctx context.Context) ([]Application, error) {
|
||||
rows, err := s.pool.Query(ctx, applicationSelect+` ORDER BY a.updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Application{}
|
||||
for rows.Next() {
|
||||
app, err := scanApplication(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, app)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
func (s *Service) GetApplication(ctx context.Context, id string) (Application, error) {
|
||||
return scanApplication(s.pool.QueryRow(ctx, applicationSelect+` WHERE a.id=$1`, id))
|
||||
}
|
||||
func (s *Service) GetPublishedApplicationByCode(ctx context.Context, code string) (Application, error) {
|
||||
return scanApplication(s.pool.QueryRow(ctx, applicationSelect+` WHERE a.code=$1 AND a.status='active' AND a.published_version IS NOT NULL`, code))
|
||||
}
|
||||
|
||||
func (s *Service) SaveApplication(ctx context.Context, app Application, actorID string, create bool) (Application, error) {
|
||||
if err := validateApplication(&app); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err := normalizeApplicationConfig(&app.DraftConfig, false); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err := s.validateApplicationRefs(ctx, app.DraftConfig); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
raw, _ := json.Marshal(app.DraftConfig)
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
if create {
|
||||
app.ID, err = newUUID()
|
||||
if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.applications(id,code,name,description,department_ids,status,draft_config,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8)`, app.ID, app.Code, app.Name, app.Description, app.DepartmentIDs, app.Status, raw, actorID)
|
||||
} else {
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.applications SET code=$2,name=$3,description=$4,department_ids=$5,status=$6,draft_config=$7,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, app.ID, app.Code, app.Name, app.Description, app.DepartmentIDs, app.Status, raw)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return Application{}, ErrNotFound
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
event := "application.updated"
|
||||
if create {
|
||||
event = "application.created"
|
||||
}
|
||||
if err = emit(ctx, tx, event, "application", app.ID, actorID, nil); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
return s.GetApplication(ctx, app.ID)
|
||||
}
|
||||
|
||||
func (s *Service) PublishApplication(ctx context.Context, id, changeNote, actorID string) (Application, error) {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
var raw []byte
|
||||
if err = tx.QueryRow(ctx, `SELECT draft_config FROM gateway.applications WHERE id=$1 FOR UPDATE`, id).Scan(&raw); errors.Is(err, pgx.ErrNoRows) {
|
||||
return Application{}, ErrNotFound
|
||||
} else if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
var config ApplicationConfig
|
||||
if err = json.Unmarshal(raw, &config); err != nil {
|
||||
return Application{}, errors.New("应用草稿配置无效")
|
||||
}
|
||||
if err = normalizeApplicationConfig(&config, true); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err = s.validateApplicationRefs(ctx, config); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
raw, _ = json.Marshal(config)
|
||||
var version int
|
||||
if err = tx.QueryRow(ctx, `SELECT coalesce(max(version),0)+1 FROM gateway.application_versions WHERE application_id=$1`, id).Scan(&version); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
versionID, idErr := newUUID()
|
||||
if idErr != nil {
|
||||
return Application{}, idErr
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `INSERT INTO gateway.application_versions(id,application_id,version,config,change_note,published_by) VALUES($1,$2,$3,$4,$5,$6)`, versionID, id, version, raw, strings.TrimSpace(changeNote), actorID); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `UPDATE gateway.applications SET published_version=$2,status='active',revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, version); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err = emit(ctx, tx, "application.published", "application", id, actorID, map[string]any{"version": version}); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
return s.GetApplication(ctx, id)
|
||||
}
|
||||
|
||||
func (s *Service) DeleteApplication(ctx context.Context, id, actorID string) error {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.applications WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "application.deleted", "application", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *Service) ListApplicationVersions(ctx context.Context, id string) ([]map[string]any, error) {
|
||||
rows, err := s.pool.Query(ctx, `SELECT id::text,version,config,change_note,created_at FROM gateway.application_versions WHERE application_id=$1 ORDER BY version DESC`, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []map[string]any{}
|
||||
for rows.Next() {
|
||||
var versionID, note string
|
||||
var version int
|
||||
var raw []byte
|
||||
var created time.Time
|
||||
if err = rows.Scan(&versionID, &version, &raw, ¬e, &created); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var config ApplicationConfig
|
||||
_ = json.Unmarshal(raw, &config)
|
||||
items = append(items, map[string]any{"id": versionID, "version": version, "config": config, "change_note": note, "created_at": created})
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) ListApplicationRuns(ctx context.Context, id string, limit int) ([]ApplicationRun, error) {
|
||||
if limit < 1 {
|
||||
limit = 50
|
||||
}
|
||||
if limit > 200 {
|
||||
limit = 200
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `SELECT id::text,application_id::text,request_id,status,error,version,latency_ms,retrieval_count,tool_count,created_at FROM gateway.application_runs WHERE application_id=$1 ORDER BY created_at DESC LIMIT $2`, id, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []ApplicationRun{}
|
||||
for rows.Next() {
|
||||
var run ApplicationRun
|
||||
if err = rows.Scan(&run.ID, &run.ApplicationID, &run.RequestID, &run.Status, &run.Error, &run.Version, &run.LatencyMS, &run.RetrievalCount, &run.ToolCount, &run.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, run)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) RollbackApplication(ctx context.Context, id string, version int, changeNote, actorID string) (Application, error) {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
var raw []byte
|
||||
if err = tx.QueryRow(ctx, `SELECT config FROM gateway.application_versions WHERE application_id=$1 AND version=$2`, id, version).Scan(&raw); errors.Is(err, pgx.ErrNoRows) {
|
||||
return Application{}, ErrNotFound
|
||||
} else if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
var config ApplicationConfig
|
||||
if err = json.Unmarshal(raw, &config); err != nil {
|
||||
return Application{}, errors.New("目标版本配置无效")
|
||||
}
|
||||
if err = normalizeApplicationConfig(&config, true); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err = s.validateApplicationRefs(ctx, config); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
var next int
|
||||
if err = tx.QueryRow(ctx, `SELECT coalesce(max(version),0)+1 FROM gateway.application_versions WHERE application_id=$1`, id).Scan(&next); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
versionID, err := newUUID()
|
||||
if err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
note := strings.TrimSpace(changeNote)
|
||||
if note == "" {
|
||||
note = fmt.Sprintf("回滚到 v%d", version)
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `INSERT INTO gateway.application_versions(id,application_id,version,config,change_note,published_by) VALUES($1,$2,$3,$4,$5,$6)`, versionID, id, next, raw, note, actorID); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `UPDATE gateway.applications SET draft_config=$2,published_version=$3,status='active',revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, raw, next); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err = emit(ctx, tx, "application.rolled_back", "application", id, actorID, map[string]any{"source_version": version, "published_version": next}); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return Application{}, err
|
||||
}
|
||||
return s.GetApplication(ctx, id)
|
||||
}
|
||||
|
||||
func (s *Service) ApplicationDependencies(ctx context.Context, assetType, assetID string) ([]map[string]any, error) {
|
||||
apps, err := s.ListApplications(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items := []map[string]any{}
|
||||
for _, app := range apps {
|
||||
stages := []struct {
|
||||
name string
|
||||
config *ApplicationConfig
|
||||
}{{"draft", &app.DraftConfig}, {"published", app.PublishedConfig}}
|
||||
for _, stage := range stages {
|
||||
if stage.config == nil {
|
||||
continue
|
||||
}
|
||||
used := assetType == "prompt" && stage.config.PromptTemplateID == assetID
|
||||
if assetType == "knowledge" {
|
||||
for _, id := range stage.config.KnowledgeBaseIDs {
|
||||
if id == assetID {
|
||||
used = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if assetType == "tool" {
|
||||
for _, id := range stage.config.ToolIDs {
|
||||
if id == assetID {
|
||||
used = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if used {
|
||||
items = append(items, map[string]any{"application_id": app.ID, "code": app.Code, "name": app.Name, "stage": stage.name, "version": app.PublishedVersion})
|
||||
}
|
||||
}
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
@@ -0,0 +1,378 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type DigitalEmployee struct {
|
||||
ID string `json:"id"`
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Persona string `json:"persona"`
|
||||
Model string `json:"model"`
|
||||
SkillIDs []string `json:"skill_ids"`
|
||||
ToolIDs []string `json:"tool_ids"`
|
||||
MCPServerIDs []string `json:"mcp_server_ids"`
|
||||
KnowledgeBaseIDs []string `json:"knowledge_base_ids"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
RetrievalTopK int `json:"retrieval_top_k"`
|
||||
MaxToolRounds int `json:"max_tool_rounds"`
|
||||
Status string `json:"status"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
CategoryName string `json:"category_name"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type DigitalEmployeeInput struct {
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Persona string `json:"persona"`
|
||||
Model string `json:"model"`
|
||||
SkillIDs []string `json:"skill_ids"`
|
||||
ToolIDs []string `json:"tool_ids"`
|
||||
MCPServerIDs []string `json:"mcp_server_ids"`
|
||||
KnowledgeBaseIDs []string `json:"knowledge_base_ids"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
RetrievalTopK int `json:"retrieval_top_k"`
|
||||
MaxToolRounds int `json:"max_tool_rounds"`
|
||||
Status string `json:"status"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
type DigitalEmployeeRun struct {
|
||||
ID string `json:"id"`
|
||||
DigitalEmployee string `json:"digital_employee"`
|
||||
RequestID string `json:"request_id"`
|
||||
Status string `json:"status"`
|
||||
Error string `json:"error"`
|
||||
LatencyMS int64 `json:"latency_ms"`
|
||||
RetrievalCount int `json:"retrieval_count"`
|
||||
ToolCount int `json:"tool_count"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// DigitalEmployeeService manages composite digital employees. A digital
|
||||
// employee bundles a persona, model, and bound skills / tools / MCP servers /
|
||||
// knowledge bases; it is the end-user consumption point of the marketplace.
|
||||
type DigitalEmployeeService struct {
|
||||
assets *Service
|
||||
skills *SkillService
|
||||
tools *ToolService
|
||||
mcpServers *MCPServerService
|
||||
}
|
||||
|
||||
func NewDigitalEmployeeService(assets *Service, skills *SkillService, tools *ToolService, mcpServers *MCPServerService) *DigitalEmployeeService {
|
||||
return &DigitalEmployeeService{assets: assets, skills: skills, tools: tools, mcpServers: mcpServers}
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) validate(ctx context.Context, input *DigitalEmployeeInput, create bool) error {
|
||||
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Description = strings.TrimSpace(input.Description)
|
||||
input.Persona = strings.TrimSpace(input.Persona)
|
||||
input.Model = strings.TrimSpace(input.Model)
|
||||
if !codeRE.MatchString(input.Code) || input.Name == "" || len(input.Name) > 128 || len(input.Description) > 4000 || len(input.Persona) > 100000 {
|
||||
return errors.New("数字员工编码、名称、描述或人设格式无效")
|
||||
}
|
||||
if input.Model == "" {
|
||||
return errors.New("必须指定模型")
|
||||
}
|
||||
if input.Temperature == 0 {
|
||||
input.Temperature = 0.7
|
||||
}
|
||||
if input.Temperature < 0 || input.Temperature > 2 {
|
||||
return errors.New("temperature 应在 0-2 之间")
|
||||
}
|
||||
if input.RetrievalTopK == 0 {
|
||||
input.RetrievalTopK = 5
|
||||
}
|
||||
if input.RetrievalTopK < 1 || input.RetrievalTopK > 50 {
|
||||
return errors.New("retrieval_top_k 应在 1-50 之间")
|
||||
}
|
||||
if input.MaxToolRounds == 0 {
|
||||
input.MaxToolRounds = 5
|
||||
}
|
||||
if input.MaxToolRounds < 1 || input.MaxToolRounds > 20 {
|
||||
return errors.New("max_tool_rounds 应在 1-20 之间")
|
||||
}
|
||||
switch input.Status {
|
||||
case "", "draft":
|
||||
input.Status = "draft"
|
||||
case "published", "archived":
|
||||
default:
|
||||
return errors.New("无效的资源状态")
|
||||
}
|
||||
var err error
|
||||
if input.SkillIDs, err = normalizeStrings(input.SkillIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.ToolIDs, err = normalizeStrings(input.ToolIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.MCPServerIDs, err = normalizeStrings(input.MCPServerIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.KnowledgeBaseIDs, err = normalizeStrings(input.KnowledgeBaseIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.CategoryID != nil && strings.TrimSpace(*input.CategoryID) == "" {
|
||||
input.CategoryID = nil
|
||||
}
|
||||
var categoryErr error
|
||||
input.CategoryID, categoryErr = validCategoryID(ctx, s.assets.pool, input.CategoryID, "digital_employee")
|
||||
if categoryErr != nil {
|
||||
return categoryErr
|
||||
}
|
||||
if input.Tags, err = normalizeStrings(input.Tags, 30); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.DepartmentIDs, err = normalizeStrings(input.DepartmentIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
return s.validateBindings(ctx, input, false)
|
||||
}
|
||||
|
||||
// validateBindings ensures every bound asset exists (and, when strict, is
|
||||
// enabled) so a published digital employee never references a missing resource.
|
||||
func (s *DigitalEmployeeService) validateBindings(ctx context.Context, input *DigitalEmployeeInput, strict bool) error {
|
||||
for _, id := range input.SkillIDs {
|
||||
skill, err := s.skills.Get(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("绑定的 Skill %s 不存在", id)
|
||||
}
|
||||
if strict && !skill.Enabled {
|
||||
return fmt.Errorf("绑定的 Skill %s 未启用", skill.Code)
|
||||
}
|
||||
}
|
||||
for _, id := range input.ToolIDs {
|
||||
tool, err := s.tools.Get(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("绑定的工具 %s 不存在", id)
|
||||
}
|
||||
if strict && !tool.Enabled {
|
||||
return fmt.Errorf("绑定的工具 %s 未启用", tool.Code)
|
||||
}
|
||||
}
|
||||
for _, id := range input.MCPServerIDs {
|
||||
server, err := s.mcpServers.Get(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("绑定的 MCP 服务器 %s 不存在", id)
|
||||
}
|
||||
if strict && !server.Enabled {
|
||||
return fmt.Errorf("绑定的 MCP 服务器 %s 未启用", server.Code)
|
||||
}
|
||||
}
|
||||
for _, id := range input.KnowledgeBaseIDs {
|
||||
var exists bool
|
||||
if err := s.assets.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.knowledge_bases WHERE id=$1)`, id).Scan(&exists); err != nil {
|
||||
return err
|
||||
}
|
||||
if !exists {
|
||||
return fmt.Errorf("绑定的知识库 %s 不存在", id)
|
||||
}
|
||||
if strict {
|
||||
var enabled bool
|
||||
if err := s.assets.pool.QueryRow(ctx, `SELECT enabled FROM gateway.knowledge_bases WHERE id=$1`, id).Scan(&enabled); err != nil {
|
||||
return err
|
||||
}
|
||||
if !enabled {
|
||||
return fmt.Errorf("绑定的知识库 %s 未启用", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const digitalEmployeeSelect = `SELECT d.id::text,d.code,d.name,d.description,d.persona,d.model,d.skill_ids,d.tool_ids,d.mcp_server_ids,d.knowledge_base_ids,d.temperature,d.retrieval_top_k,d.max_tool_rounds,d.status,d.category_id::text,coalesce(c.name,''),d.tags,d.department_ids::text[],d.enabled,d.revision,d.created_at,d.updated_at FROM gateway.digital_employees d LEFT JOIN gateway.marketplace_categories c ON c.id=d.category_id`
|
||||
|
||||
func scanDigitalEmployee(row pgx.Row) (DigitalEmployee, error) {
|
||||
var d DigitalEmployee
|
||||
err := row.Scan(&d.ID, &d.Code, &d.Name, &d.Description, &d.Persona, &d.Model, &d.SkillIDs, &d.ToolIDs, &d.MCPServerIDs, &d.KnowledgeBaseIDs, &d.Temperature, &d.RetrievalTopK, &d.MaxToolRounds, &d.Status, &d.CategoryID, &d.CategoryName, &d.Tags, &d.DepartmentIDs, &d.Enabled, &d.Revision, &d.CreatedAt, &d.UpdatedAt)
|
||||
return d, mapNotFound(err)
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) List(ctx context.Context) ([]DigitalEmployee, error) {
|
||||
rows, err := s.assets.pool.Query(ctx, digitalEmployeeSelect+` ORDER BY d.updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []DigitalEmployee{}
|
||||
for rows.Next() {
|
||||
item, err := scanDigitalEmployee(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) Get(ctx context.Context, id string) (DigitalEmployee, error) {
|
||||
return scanDigitalEmployee(s.assets.pool.QueryRow(ctx, digitalEmployeeSelect+` WHERE d.id=$1`, id))
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) GetByCode(ctx context.Context, code string) (DigitalEmployee, error) {
|
||||
return scanDigitalEmployee(s.assets.pool.QueryRow(ctx, digitalEmployeeSelect+` WHERE d.code=$1`, code))
|
||||
}
|
||||
|
||||
// GetPublishedByCode returns a published, enabled digital employee by code.
|
||||
func (s *DigitalEmployeeService) GetPublishedByCode(ctx context.Context, code string) (DigitalEmployee, error) {
|
||||
return scanDigitalEmployee(s.assets.pool.QueryRow(ctx, digitalEmployeeSelect+` WHERE d.code=$1 AND d.status='published' AND d.enabled`, code))
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) Save(ctx context.Context, id string, input DigitalEmployeeInput, actorID string, create bool) (DigitalEmployee, error) {
|
||||
if err := s.validate(ctx, &input, create); err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
if create {
|
||||
id, err = newUUID()
|
||||
if err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.digital_employees(id,code,name,description,persona,model,skill_ids,tool_ids,mcp_server_ids,knowledge_base_ids,temperature,retrieval_top_k,max_tool_rounds,status,category_id,tags,department_ids,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19)`, id, input.Code, input.Name, input.Description, input.Persona, input.Model, input.SkillIDs, input.ToolIDs, input.MCPServerIDs, input.KnowledgeBaseIDs, input.Temperature, input.RetrievalTopK, input.MaxToolRounds, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled, actorID)
|
||||
} else {
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.digital_employees SET code=$2,name=$3,description=$4,persona=$5,model=$6,skill_ids=$7,tool_ids=$8,mcp_server_ids=$9,knowledge_base_ids=$10,temperature=$11,retrieval_top_k=$12,max_tool_rounds=$13,status=$14,category_id=$15,tags=$16,department_ids=$17,enabled=$18,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.Persona, input.Model, input.SkillIDs, input.ToolIDs, input.MCPServerIDs, input.KnowledgeBaseIDs, input.Temperature, input.RetrievalTopK, input.MaxToolRounds, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return DigitalEmployee{}, ErrNotFound
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
event := "digital_employee.updated"
|
||||
if create {
|
||||
event = "digital_employee.created"
|
||||
}
|
||||
if err = emit(ctx, tx, event, "digital_employee", id, actorID, nil); err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
// Publish revalidates all bindings (must exist and be enabled) and flips the
|
||||
// status to published so the resource appears in the marketplace.
|
||||
func (s *DigitalEmployeeService) Publish(ctx context.Context, id, actorID string) (DigitalEmployee, error) {
|
||||
current, err := s.Get(ctx, id)
|
||||
if err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
input := DigitalEmployeeInput{
|
||||
Code: current.Code, Name: current.Name, Description: current.Description,
|
||||
Persona: current.Persona, Model: current.Model,
|
||||
SkillIDs: current.SkillIDs, ToolIDs: current.ToolIDs, MCPServerIDs: current.MCPServerIDs,
|
||||
KnowledgeBaseIDs: current.KnowledgeBaseIDs, Temperature: current.Temperature,
|
||||
RetrievalTopK: current.RetrievalTopK, MaxToolRounds: current.MaxToolRounds,
|
||||
Status: "published", CategoryID: current.CategoryID, Tags: current.Tags,
|
||||
DepartmentIDs: current.DepartmentIDs, Enabled: current.Enabled,
|
||||
}
|
||||
if err = s.validateBindings(ctx, &input, true); err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
if err = s.setStatus(ctx, id, actorID, "published"); err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
// Archive unpublishes a digital employee from the marketplace without deleting
|
||||
// its configuration.
|
||||
func (s *DigitalEmployeeService) Archive(ctx context.Context, id, actorID string) (DigitalEmployee, error) {
|
||||
if err := s.setStatus(ctx, id, actorID, "archived"); err != nil {
|
||||
return DigitalEmployee{}, err
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) setStatus(ctx context.Context, id, actorID, status string) error {
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `UPDATE gateway.digital_employees SET status=$2,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, status)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "digital_employee."+status, "digital_employee", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) 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.marketplace_installations WHERE resource_type='digital_employee' AND resource_id=$1)`, id).Scan(&used); err != nil {
|
||||
return err
|
||||
}
|
||||
if used {
|
||||
return ErrConflict
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.digital_employees WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "digital_employee.deleted", "digital_employee", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *DigitalEmployeeService) ListRuns(ctx context.Context, id string, limit int) ([]DigitalEmployeeRun, error) {
|
||||
if limit < 1 {
|
||||
limit = 50
|
||||
}
|
||||
if limit > 200 {
|
||||
limit = 200
|
||||
}
|
||||
rows, err := s.assets.pool.Query(ctx, `SELECT id::text,(SELECT code FROM gateway.digital_employees WHERE id=r.digital_employee_id),request_id,status,error,latency_ms,retrieval_count,tool_count,created_at FROM gateway.digital_employee_runs r WHERE digital_employee_id=$1 ORDER BY created_at DESC LIMIT $2`, id, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []DigitalEmployeeRun{}
|
||||
for rows.Next() {
|
||||
var run DigitalEmployeeRun
|
||||
if err = rows.Scan(&run.ID, &run.DigitalEmployee, &run.RequestID, &run.Status, &run.Error, &run.LatencyMS, &run.RetrievalCount, &run.ToolCount, &run.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, run)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,448 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type Retriever interface {
|
||||
Search(context.Context, string, string, int) ([]SearchHit, error)
|
||||
}
|
||||
|
||||
type PostgreSQLRetriever struct {
|
||||
pool interface {
|
||||
Query(context.Context, string, ...any) (pgx.Rows, error)
|
||||
}
|
||||
}
|
||||
|
||||
func NewPostgreSQLRetriever(service *Service) *PostgreSQLRetriever {
|
||||
return &PostgreSQLRetriever{pool: service.pool}
|
||||
}
|
||||
|
||||
func (r *PostgreSQLRetriever) Search(ctx context.Context, knowledgeBaseID, query string, topK int) ([]SearchHit, error) {
|
||||
query = strings.TrimSpace(query)
|
||||
if query == "" {
|
||||
return nil, errors.New("检索词不能为空")
|
||||
}
|
||||
if len(query) > 6000 {
|
||||
return nil, errors.New("检索词过长")
|
||||
}
|
||||
if topK < 1 {
|
||||
topK = 4
|
||||
}
|
||||
if topK > 20 {
|
||||
topK = 20
|
||||
}
|
||||
tokens := searchTokens(query)
|
||||
rows, err := r.pool.Query(ctx, `WITH q AS (SELECT plainto_tsquery('simple',$2) AS tsq,lower($2) AS raw), tokens AS (SELECT unnest($4::text[]) AS token)
|
||||
SELECT c.id::text,c.document_id::text,d.title,c.chunk_index,c.content,
|
||||
greatest(ts_rank_cd(c.search_vector,q.tsq),CASE WHEN strpos(lower(c.content),q.raw)>0 THEN 1.0 ELSE 0.0 END,
|
||||
(SELECT count(*)::float8/greatest(array_length($4::text[],1),1) FROM tokens WHERE strpos(lower(c.content),token)>0))::float8 AS score
|
||||
FROM gateway.knowledge_chunks c JOIN gateway.knowledge_documents d ON d.id=c.document_id JOIN gateway.knowledge_bases k ON k.id=c.knowledge_base_id CROSS JOIN q
|
||||
WHERE c.knowledge_base_id=$1 AND k.enabled AND d.status='ready' AND (c.search_vector @@ q.tsq OR strpos(lower(c.content),q.raw)>0 OR EXISTS(SELECT 1 FROM tokens WHERE strpos(lower(c.content),token)>0))
|
||||
ORDER BY score DESC,c.document_id,c.chunk_index LIMIT $3`, knowledgeBaseID, query, topK, tokens)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
hits := []SearchHit{}
|
||||
for rows.Next() {
|
||||
var h SearchHit
|
||||
if err = rows.Scan(&h.ChunkID, &h.DocumentID, &h.DocumentTitle, &h.ChunkIndex, &h.Content, &h.Score); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hits = append(hits, h)
|
||||
}
|
||||
return hits, rows.Err()
|
||||
}
|
||||
|
||||
func searchTokens(query string) []string {
|
||||
query = strings.ToLower(strings.TrimSpace(query))
|
||||
seen := map[string]struct{}{}
|
||||
tokens := make([]string, 0, 32)
|
||||
add := func(value string) {
|
||||
value = strings.TrimSpace(value)
|
||||
if utf8.RuneCountInString(value) < 2 {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[value]; ok {
|
||||
return
|
||||
}
|
||||
seen[value] = struct{}{}
|
||||
tokens = append(tokens, value)
|
||||
}
|
||||
words := strings.FieldsFunc(query, func(value rune) bool { return !unicode.IsLetter(value) && !unicode.IsNumber(value) })
|
||||
for _, word := range words {
|
||||
runes := []rune(word)
|
||||
hasCJK := false
|
||||
for _, value := range runes {
|
||||
if unicode.In(value, unicode.Han) {
|
||||
hasCJK = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasCJK {
|
||||
add(word)
|
||||
continue
|
||||
}
|
||||
for index := 0; index+1 < len(runes); index++ {
|
||||
add(string(runes[index : index+2]))
|
||||
}
|
||||
}
|
||||
if len(tokens) > 64 {
|
||||
tokens = tokens[:64]
|
||||
}
|
||||
sort.Strings(tokens)
|
||||
return tokens
|
||||
}
|
||||
|
||||
func validateKnowledgeBase(k *KnowledgeBase) error {
|
||||
k.Name = strings.TrimSpace(k.Name)
|
||||
k.Description = strings.TrimSpace(k.Description)
|
||||
if k.Name == "" || len(k.Name) > 128 || len(k.Description) > 4000 {
|
||||
return errors.New("知识库名称或描述格式无效")
|
||||
}
|
||||
if k.RetrievalMode == "" {
|
||||
k.RetrievalMode = "postgres_fts"
|
||||
}
|
||||
if k.RetrievalMode != "postgres_fts" {
|
||||
return errors.New("基线仅支持 postgres_fts 检索器")
|
||||
}
|
||||
if k.ChunkSize == 0 {
|
||||
k.ChunkSize = 800
|
||||
}
|
||||
if k.ChunkSize < 200 || k.ChunkSize > 8000 {
|
||||
return errors.New("chunk_size 应在 200-8000 之间")
|
||||
}
|
||||
if k.ChunkOverlap < 0 || k.ChunkOverlap > k.ChunkSize/2 {
|
||||
return errors.New("chunk_overlap 应在 0 到 chunk_size 一半之间")
|
||||
}
|
||||
var err error
|
||||
k.DepartmentIDs, err = normalizeStrings(k.DepartmentIDs, 100)
|
||||
return err
|
||||
}
|
||||
|
||||
const knowledgeBaseSelect = `SELECT k.id::text,k.name,k.description,k.retrieval_mode,k.chunk_size,k.chunk_overlap,k.department_ids::text[],k.enabled,k.revision,
|
||||
count(DISTINCT d.id)::int,count(c.id)::int,k.created_at,k.updated_at
|
||||
FROM gateway.knowledge_bases k LEFT JOIN gateway.knowledge_documents d ON d.knowledge_base_id=k.id LEFT JOIN gateway.knowledge_chunks c ON c.document_id=d.id`
|
||||
|
||||
func scanKnowledgeBase(row pgx.Row) (KnowledgeBase, error) {
|
||||
var k KnowledgeBase
|
||||
err := row.Scan(&k.ID, &k.Name, &k.Description, &k.RetrievalMode, &k.ChunkSize, &k.ChunkOverlap, &k.DepartmentIDs, &k.Enabled, &k.Revision, &k.DocumentCount, &k.ChunkCount, &k.CreatedAt, &k.UpdatedAt)
|
||||
return k, mapNotFound(err)
|
||||
}
|
||||
|
||||
func (s *Service) ListKnowledgeBases(ctx context.Context) ([]KnowledgeBase, error) {
|
||||
rows, err := s.pool.Query(ctx, knowledgeBaseSelect+` GROUP BY k.id ORDER BY k.updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []KnowledgeBase{}
|
||||
for rows.Next() {
|
||||
k, err := scanKnowledgeBase(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, k)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
func (s *Service) GetKnowledgeBase(ctx context.Context, id string) (KnowledgeBase, error) {
|
||||
return scanKnowledgeBase(s.pool.QueryRow(ctx, knowledgeBaseSelect+` WHERE k.id=$1 GROUP BY k.id`, id))
|
||||
}
|
||||
|
||||
func (s *Service) SaveKnowledgeBase(ctx context.Context, k KnowledgeBase, actorID string, create bool) (KnowledgeBase, error) {
|
||||
if err := validateKnowledgeBase(&k); err != nil {
|
||||
return KnowledgeBase{}, err
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return KnowledgeBase{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
if create {
|
||||
k.ID, err = newUUID()
|
||||
if err != nil {
|
||||
return KnowledgeBase{}, err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.knowledge_bases(id,name,description,retrieval_mode,chunk_size,chunk_overlap,department_ids,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9)`, k.ID, k.Name, k.Description, k.RetrievalMode, k.ChunkSize, k.ChunkOverlap, k.DepartmentIDs, k.Enabled, actorID)
|
||||
} else {
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.knowledge_bases SET name=$2,description=$3,retrieval_mode=$4,chunk_size=$5,chunk_overlap=$6,department_ids=$7,enabled=$8,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, k.ID, k.Name, k.Description, k.RetrievalMode, k.ChunkSize, k.ChunkOverlap, k.DepartmentIDs, k.Enabled)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return KnowledgeBase{}, ErrNotFound
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return KnowledgeBase{}, err
|
||||
}
|
||||
event := "knowledge_base.updated"
|
||||
if create {
|
||||
event = "knowledge_base.created"
|
||||
}
|
||||
if err = emit(ctx, tx, event, "knowledge_base", k.ID, actorID, nil); err != nil {
|
||||
return KnowledgeBase{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return KnowledgeBase{}, err
|
||||
}
|
||||
return s.GetKnowledgeBase(ctx, k.ID)
|
||||
}
|
||||
|
||||
func (s *Service) DeleteKnowledgeBase(ctx context.Context, id, actorID string) error {
|
||||
tx, err := s.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->'knowledge_base_ids' ? $1 UNION ALL SELECT 1 FROM gateway.application_versions WHERE config->'knowledge_base_ids' ? $1)`, id).Scan(&used); err != nil {
|
||||
return err
|
||||
}
|
||||
if used {
|
||||
return ErrConflict
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.knowledge_bases WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "knowledge_base.deleted", "knowledge_base", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func ChunkText(text string, size, overlap int) []string {
|
||||
if size < 200 {
|
||||
size = 800
|
||||
}
|
||||
if overlap < 0 {
|
||||
overlap = 0
|
||||
}
|
||||
if overlap > size/2 {
|
||||
overlap = size / 2
|
||||
}
|
||||
paragraphs := regexp.MustCompile(`\n\s*\n`).Split(strings.TrimSpace(text), -1)
|
||||
chunks := []string{}
|
||||
current := []rune{}
|
||||
appendCurrent := func() {
|
||||
value := strings.TrimSpace(string(current))
|
||||
if value != "" {
|
||||
chunks = append(chunks, value)
|
||||
}
|
||||
current = nil
|
||||
}
|
||||
for _, paragraph := range paragraphs {
|
||||
runes := []rune(strings.TrimSpace(paragraph))
|
||||
if len(runes) == 0 {
|
||||
continue
|
||||
}
|
||||
if len(runes) > size {
|
||||
appendCurrent()
|
||||
step := size - overlap
|
||||
if step < 1 {
|
||||
step = 1
|
||||
}
|
||||
for start := 0; start < len(runes); start += step {
|
||||
end := start + size
|
||||
if end > len(runes) {
|
||||
end = len(runes)
|
||||
}
|
||||
chunks = append(chunks, string(runes[start:end]))
|
||||
if end == len(runes) {
|
||||
break
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
separator := 0
|
||||
if len(current) > 0 {
|
||||
separator = 2
|
||||
}
|
||||
if len(current)+separator+len(runes) <= size {
|
||||
if separator > 0 {
|
||||
current = append(current, '\n', '\n')
|
||||
}
|
||||
current = append(current, runes...)
|
||||
continue
|
||||
}
|
||||
previous := append([]rune(nil), current...)
|
||||
appendCurrent()
|
||||
if overlap > 0 && len(previous) > 0 {
|
||||
start := len(previous) - overlap
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
current = append(current, previous[start:]...)
|
||||
current = append(current, '\n', '\n')
|
||||
}
|
||||
current = append(current, runes...)
|
||||
}
|
||||
appendCurrent()
|
||||
return chunks
|
||||
}
|
||||
|
||||
func (s *Service) AddKnowledgeDocument(ctx context.Context, kbID, title, sourceType, sourceURI, content, actorID string) (KnowledgeDocument, error) {
|
||||
title = strings.TrimSpace(title)
|
||||
sourceType = strings.TrimSpace(sourceType)
|
||||
sourceURI = strings.TrimSpace(sourceURI)
|
||||
content = strings.TrimSpace(content)
|
||||
if title == "" || len(title) > 256 {
|
||||
return KnowledgeDocument{}, errors.New("文档标题格式无效")
|
||||
}
|
||||
if sourceType == "" {
|
||||
sourceType = "text"
|
||||
}
|
||||
if sourceType != "text" && sourceType != "url" && sourceType != "import" {
|
||||
return KnowledgeDocument{}, errors.New("source_type 无效")
|
||||
}
|
||||
if content == "" || len([]byte(content)) > 2<<20 {
|
||||
return KnowledgeDocument{}, errors.New("文档正文不能为空且最多 2 MiB")
|
||||
}
|
||||
kb, err := s.GetKnowledgeBase(ctx, kbID)
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
chunks := ChunkText(content, kb.ChunkSize, kb.ChunkOverlap)
|
||||
if len(chunks) == 0 {
|
||||
return KnowledgeDocument{}, errors.New("文档没有可入库内容")
|
||||
}
|
||||
digest := sha256.Sum256([]byte(content))
|
||||
hash := hex.EncodeToString(digest[:])
|
||||
docID, err := newUUID()
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
var doc KnowledgeDocument
|
||||
err = tx.QueryRow(ctx, `INSERT INTO gateway.knowledge_documents(id,knowledge_base_id,title,source_type,source_uri,content,content_sha256,chunk_count,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id::text,knowledge_base_id::text,title,source_type,source_uri,content_sha256,status,status_message,char_length(content),chunk_count,created_at,updated_at`, docID, kbID, title, sourceType, sourceURI, content, hash, len(chunks), actorID).Scan(&doc.ID, &doc.KnowledgeBaseID, &doc.Title, &doc.SourceType, &doc.SourceURI, &doc.ContentSHA256, &doc.Status, &doc.StatusMessage, &doc.CharCount, &doc.ChunkCount, &doc.CreatedAt, &doc.UpdatedAt)
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
chunkRows := make([][]any, 0, len(chunks))
|
||||
for index, chunk := range chunks {
|
||||
chunkID, idErr := newUUID()
|
||||
if idErr != nil {
|
||||
return KnowledgeDocument{}, idErr
|
||||
}
|
||||
chunkRows = append(chunkRows, []any{chunkID, kbID, docID, index, chunk})
|
||||
}
|
||||
// Bulk-copy all chunks in one statement instead of one INSERT per chunk;
|
||||
// a 2 MiB document can split into thousands of chunks.
|
||||
if _, err = tx.CopyFrom(ctx, pgx.Identifier{"gateway", "knowledge_chunks"}, []string{"id", "knowledge_base_id", "document_id", "chunk_index", "content"}, pgx.CopyFromRows(chunkRows)); err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
if err = emit(ctx, tx, "knowledge_document.ready", "knowledge_document", docID, actorID, map[string]any{"knowledge_base_id": kbID, "chunk_count": len(chunks)}); err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
func (s *Service) ListKnowledgeDocuments(ctx context.Context, kbID string) ([]KnowledgeDocument, error) {
|
||||
rows, err := s.pool.Query(ctx, `SELECT id::text,knowledge_base_id::text,title,source_type,source_uri,content_sha256,status,status_message,char_length(content),chunk_count,created_at,updated_at FROM gateway.knowledge_documents WHERE knowledge_base_id=$1 ORDER BY created_at DESC`, kbID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []KnowledgeDocument{}
|
||||
for rows.Next() {
|
||||
var d KnowledgeDocument
|
||||
if err = rows.Scan(&d.ID, &d.KnowledgeBaseID, &d.Title, &d.SourceType, &d.SourceURI, &d.ContentSHA256, &d.Status, &d.StatusMessage, &d.CharCount, &d.ChunkCount, &d.CreatedAt, &d.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, d)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) DeleteKnowledgeDocument(ctx context.Context, kbID, id, actorID string) error {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.knowledge_documents WHERE id=$1 AND knowledge_base_id=$2`, id, kbID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "knowledge_document.deleted", "knowledge_document", id, actorID, map[string]any{"knowledge_base_id": kbID}); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *Service) ReprocessKnowledgeDocument(ctx context.Context, kbID, id, actorID string) (KnowledgeDocument, error) {
|
||||
kb, err := s.GetKnowledgeBase(ctx, kbID)
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
var content string
|
||||
var doc KnowledgeDocument
|
||||
err = s.pool.QueryRow(ctx, `SELECT id::text,knowledge_base_id::text,title,source_type,source_uri,content,content_sha256,status,status_message,char_length(content),chunk_count,created_at,updated_at FROM gateway.knowledge_documents WHERE id=$1 AND knowledge_base_id=$2`, id, kbID).Scan(&doc.ID, &doc.KnowledgeBaseID, &doc.Title, &doc.SourceType, &doc.SourceURI, &content, &doc.ContentSHA256, &doc.Status, &doc.StatusMessage, &doc.CharCount, &doc.ChunkCount, &doc.CreatedAt, &doc.UpdatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return KnowledgeDocument{}, ErrNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
chunks := ChunkText(content, kb.ChunkSize, kb.ChunkOverlap)
|
||||
if len(chunks) == 0 {
|
||||
return KnowledgeDocument{}, errors.New("文档没有可入库内容")
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
if _, err = tx.Exec(ctx, `DELETE FROM gateway.knowledge_chunks WHERE document_id=$1`, id); err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
for index, chunk := range chunks {
|
||||
chunkID, idErr := newUUID()
|
||||
if idErr != nil {
|
||||
return KnowledgeDocument{}, idErr
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `INSERT INTO gateway.knowledge_chunks(id,knowledge_base_id,document_id,chunk_index,content) VALUES($1,$2,$3,$4,$5)`, chunkID, kbID, id, index, chunk); err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
}
|
||||
err = tx.QueryRow(ctx, `UPDATE gateway.knowledge_documents SET status='ready',status_message='',chunk_count=$3,updated_at=clock_timestamp() WHERE id=$1 AND knowledge_base_id=$2 RETURNING updated_at`, id, kbID, len(chunks)).Scan(&doc.UpdatedAt)
|
||||
if err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
doc.Status, doc.StatusMessage, doc.ChunkCount = "ready", "", len(chunks)
|
||||
if err = emit(ctx, tx, "knowledge_document.reprocessed", "knowledge_document", id, actorID, map[string]any{"knowledge_base_id": kbID, "chunk_count": len(chunks)}); err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return KnowledgeDocument{}, err
|
||||
}
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
func RuneCount(value string) int { return utf8.RuneCountInString(value) }
|
||||
|
||||
var _ Retriever = (*PostgreSQLRetriever)(nil)
|
||||
@@ -0,0 +1,444 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// MarketItem is the lightweight unified catalog row for a published resource,
|
||||
// regardless of which of the three resource tables it lives in.
|
||||
type MarketItem struct {
|
||||
Type string `json:"type"`
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
CategoryName string `json:"category_name"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Category is a shared marketplace category, optionally scoped to one resource
|
||||
// type (empty resource_type = global).
|
||||
type Category struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
ResourceType string `json:"resource_type"`
|
||||
SortOrder int `json:"sort_order"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// MarketplaceService provides the shared cross-resource surface: categories,
|
||||
// the unified catalog, and portal workspace installations.
|
||||
type MarketplaceService struct {
|
||||
assets *Service
|
||||
mcpServers *MCPServerService
|
||||
skills *SkillService
|
||||
employees *DigitalEmployeeService
|
||||
}
|
||||
|
||||
func NewMarketplaceService(assets *Service, mcpServers *MCPServerService, skills *SkillService, employees *DigitalEmployeeService) *MarketplaceService {
|
||||
return &MarketplaceService{assets: assets, mcpServers: mcpServers, skills: skills, employees: employees}
|
||||
}
|
||||
|
||||
// --- 分类 ---
|
||||
|
||||
func (s *MarketplaceService) ListCategories(ctx context.Context, resourceType string) ([]Category, error) {
|
||||
rows, err := s.assets.pool.Query(ctx, `SELECT id::text,name,description,resource_type,sort_order,created_at FROM gateway.marketplace_categories WHERE resource_type=$1 OR resource_type='' ORDER BY sort_order,created_at`, resourceType)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Category{}
|
||||
for rows.Next() {
|
||||
var c Category
|
||||
if err = rows.Scan(&c.ID, &c.Name, &c.Description, &c.ResourceType, &c.SortOrder, &c.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, c)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) CreateCategory(ctx context.Context, input Category, actorID string) (Category, error) {
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Description = strings.TrimSpace(input.Description)
|
||||
if input.Name == "" || len(input.Name) > 64 || len(input.Description) > 1000 {
|
||||
return Category{}, errors.New("分类名称或描述格式无效")
|
||||
}
|
||||
switch input.ResourceType {
|
||||
case "", "mcp_server", "skill", "digital_employee":
|
||||
default:
|
||||
return Category{}, errors.New("无效的分类资源类型")
|
||||
}
|
||||
if input.SortOrder < 0 {
|
||||
input.SortOrder = 0
|
||||
}
|
||||
id, err := newUUID()
|
||||
if err != nil {
|
||||
return Category{}, err
|
||||
}
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Category{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.marketplace_categories(id,name,description,resource_type,sort_order,created_by) VALUES($1,$2,$3,$4,$5,$6)`, id, input.Name, input.Description, input.ResourceType, input.SortOrder, actorID)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "duplicate") {
|
||||
return Category{}, ErrConflict
|
||||
}
|
||||
return Category{}, err
|
||||
}
|
||||
if err = emit(ctx, tx, "marketplace_category.created", "marketplace_category", id, actorID, nil); err != nil {
|
||||
return Category{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return Category{}, err
|
||||
}
|
||||
input.ID = id
|
||||
input.CreatedAt = time.Now()
|
||||
return input, nil
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) DeleteCategory(ctx context.Context, id, actorID string) error {
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.marketplace_categories WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "marketplace_category.deleted", "marketplace_category", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
// validCategoryID resolves and verifies an optional category reference. A nil
|
||||
// or blank id yields a nil category; a non-empty id must exist and match the
|
||||
// resource type scope (global categories apply to every type).
|
||||
func validCategoryID(ctx context.Context, pool *pgxpool.Pool, id *string, resourceType string) (*string, error) {
|
||||
if id == nil || strings.TrimSpace(*id) == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var exists bool
|
||||
err := pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.marketplace_categories WHERE id=$1 AND (resource_type='' OR resource_type=$2))`, *id, resourceType).Scan(&exists)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !exists {
|
||||
return nil, errors.New("分类不存在或与资源类型不匹配")
|
||||
}
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// --- 统一目录 ---
|
||||
|
||||
// Catalog returns published, enabled resources across all three types, merged
|
||||
// and sorted by update time. Filters are optional and compose with AND.
|
||||
func (s *MarketplaceService) Catalog(ctx context.Context, resourceType, categoryID, tag, query string, limit int) ([]MarketItem, error) {
|
||||
if limit < 1 {
|
||||
limit = 100
|
||||
}
|
||||
if limit > 200 {
|
||||
limit = 200
|
||||
}
|
||||
items := make([]MarketItem, 0, limit)
|
||||
types := []string{"mcp_server", "skill", "digital_employee"}
|
||||
if resourceType != "" {
|
||||
types = []string{resourceType}
|
||||
}
|
||||
for _, typ := range types {
|
||||
// Filters are built per table alias so name/description don't clash
|
||||
// with the LEFT JOINed category columns.
|
||||
where, args := marketFilters(aliasFor(typ), categoryID, tag, query)
|
||||
if err := s.catalogTable(ctx, typ, where, args, limit, &items); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
sort.SliceStable(items, func(i, j int) bool { return items[i].UpdatedAt.After(items[j].UpdatedAt) })
|
||||
if len(items) > limit {
|
||||
items = items[:limit]
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func aliasFor(resourceType string) string {
|
||||
switch resourceType {
|
||||
case "mcp_server":
|
||||
return "m"
|
||||
case "skill":
|
||||
return "s"
|
||||
default:
|
||||
return "d"
|
||||
}
|
||||
}
|
||||
|
||||
func marketFilters(alias, categoryID, tag, query string) (string, []any) {
|
||||
clauses := []string{}
|
||||
args := []any{}
|
||||
if categoryID != "" {
|
||||
args = append(args, categoryID)
|
||||
clauses = append(clauses, fmt.Sprintf("%s.category_id=$%d", alias, len(args)))
|
||||
}
|
||||
if tag != "" {
|
||||
args = append(args, tag)
|
||||
clauses = append(clauses, fmt.Sprintf("$%d = ANY(%s.tags)", len(args), alias))
|
||||
}
|
||||
if query != "" {
|
||||
args = append(args, "%"+query+"%")
|
||||
clauses = append(clauses, fmt.Sprintf("(%s.name ILIKE $%d OR %s.code ILIKE $%d OR %s.description ILIKE $%d)", alias, len(args), alias, len(args), alias, len(args)))
|
||||
}
|
||||
where := "status='published' AND enabled"
|
||||
if len(clauses) > 0 {
|
||||
where += " AND " + strings.Join(clauses, " AND ")
|
||||
}
|
||||
return where, args
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) catalogTable(ctx context.Context, typ, where string, args []any, limit int, out *[]MarketItem) error {
|
||||
// Filters occupy $1..$N in the WHERE clause; the LIMIT placeholder must be
|
||||
// the next position (N+1), not a hardcoded $1 which would collide with a
|
||||
// category filter arg and land a uuid in LIMIT.
|
||||
limitParam := fmt.Sprintf("$%d", len(args)+1)
|
||||
var query string
|
||||
switch typ {
|
||||
case "mcp_server":
|
||||
query = `SELECT 'mcp_server',m.code,m.name,m.description,m.category_id::text,coalesce(c.name,''),m.tags,m.department_ids::text[],m.updated_at FROM gateway.mcp_servers m LEFT JOIN gateway.marketplace_categories c ON c.id=m.category_id WHERE ` + where + ` ORDER BY m.updated_at DESC LIMIT ` + limitParam
|
||||
case "skill":
|
||||
query = `SELECT 'skill',s.code,s.name,s.description,s.category_id::text,coalesce(c.name,''),s.tags,s.department_ids::text[],s.updated_at FROM gateway.skills s LEFT JOIN gateway.marketplace_categories c ON c.id=s.category_id WHERE ` + where + ` ORDER BY s.updated_at DESC LIMIT ` + limitParam
|
||||
case "digital_employee":
|
||||
query = `SELECT 'digital_employee',d.code,d.name,d.description,d.category_id::text,coalesce(c.name,''),d.tags,d.department_ids::text[],d.updated_at FROM gateway.digital_employees d LEFT JOIN gateway.marketplace_categories c ON c.id=d.category_id WHERE ` + where + ` ORDER BY d.updated_at DESC LIMIT ` + limitParam
|
||||
default:
|
||||
return errors.New("未知的资源类型")
|
||||
}
|
||||
queryArgs := append(append([]any{}, args...), limit)
|
||||
rows, err := s.assets.pool.Query(ctx, query, queryArgs...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var item MarketItem
|
||||
if err = rows.Scan(&item.Type, &item.Code, &item.Name, &item.Description, &item.CategoryID, &item.CategoryName, &item.Tags, &item.DepartmentIDs, &item.UpdatedAt); err != nil {
|
||||
return err
|
||||
}
|
||||
*out = append(*out, item)
|
||||
}
|
||||
return rows.Err()
|
||||
}
|
||||
|
||||
// Detail returns the full resource payload for a published item by type+code.
|
||||
// MCP servers are returned without their encrypted headers; the consumer calls
|
||||
// a dedicated endpoint to test connectivity.
|
||||
func (s *MarketplaceService) Detail(ctx context.Context, resourceType, code string) (MarketItem, json.RawMessage, error) {
|
||||
var item MarketItem
|
||||
var detail any
|
||||
var err error
|
||||
switch resourceType {
|
||||
case "mcp_server":
|
||||
server, getErr := s.mcpServers.GetPublishedByCode(ctx, code)
|
||||
if getErr != nil {
|
||||
return item, nil, getErr
|
||||
}
|
||||
item = MarketItem{Type: "mcp_server", Code: server.Code, Name: server.Name, Description: server.Description, CategoryID: server.CategoryID, CategoryName: server.CategoryName, Tags: server.Tags, DepartmentIDs: server.DepartmentIDs, UpdatedAt: server.UpdatedAt}
|
||||
detail = server
|
||||
case "skill":
|
||||
skill, getErr := s.skills.GetPublishedByCode(ctx, code)
|
||||
if getErr != nil {
|
||||
return item, nil, getErr
|
||||
}
|
||||
item = MarketItem{Type: "skill", Code: skill.Code, Name: skill.Name, Description: skill.Description, CategoryID: skill.CategoryID, CategoryName: skill.CategoryName, Tags: skill.Tags, DepartmentIDs: skill.DepartmentIDs, UpdatedAt: skill.UpdatedAt}
|
||||
detail = skill
|
||||
case "digital_employee":
|
||||
employee, getErr := s.employees.GetPublishedByCode(ctx, code)
|
||||
if getErr != nil {
|
||||
return item, nil, getErr
|
||||
}
|
||||
item = MarketItem{Type: "digital_employee", Code: employee.Code, Name: employee.Name, Description: employee.Description, CategoryID: employee.CategoryID, CategoryName: employee.CategoryName, Tags: employee.Tags, DepartmentIDs: employee.DepartmentIDs, UpdatedAt: employee.UpdatedAt}
|
||||
detail = employee
|
||||
default:
|
||||
return item, nil, errors.New("未知的资源类型")
|
||||
}
|
||||
raw, err := json.Marshal(detail)
|
||||
if err != nil {
|
||||
return item, nil, err
|
||||
}
|
||||
return item, raw, nil
|
||||
}
|
||||
|
||||
// --- 安装(工作区绑定 + 权限) ---
|
||||
|
||||
// Install records a portal user's workspace binding to a published resource.
|
||||
// It is the permission grant that lets a cross-department user invoke a
|
||||
// resource that would otherwise be invisible to them.
|
||||
func (s *MarketplaceService) Install(ctx context.Context, resourceType, code, portalUserID string) (bool, error) {
|
||||
resourceID, err := s.publishedResourceID(ctx, resourceType, code)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
id, err := newUUID()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `INSERT INTO gateway.marketplace_installations(id,resource_type,resource_id,portal_user_id) VALUES($1,$2,$3,$4) ON CONFLICT(resource_type,resource_id,portal_user_id) DO NOTHING`, id, resourceType, resourceID, portalUserID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
created := tag.RowsAffected() == 1
|
||||
// Idempotent re-install: only emit the audit event when a row was created.
|
||||
if created {
|
||||
if err = emit(ctx, tx, "marketplace.installed", resourceType, resourceID, "", map[string]any{"code": code, "portal_user_id": portalUserID}); err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return created, nil
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) Uninstall(ctx context.Context, resourceType, code, portalUserID string) error {
|
||||
resourceID, err := s.publishedResourceID(ctx, resourceType, code)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.marketplace_installations WHERE resource_type=$1 AND resource_id=$2 AND portal_user_id=$3`, resourceType, resourceID, portalUserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "marketplace.uninstalled", resourceType, resourceID, "", map[string]any{"code": code, "portal_user_id": portalUserID}); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
// ListInstalled returns the codes the portal user has installed, grouped by
|
||||
// resource type, joined with the live resource name for display.
|
||||
func (s *MarketplaceService) ListInstalled(ctx context.Context, portalUserID string) ([]MarketItem, error) {
|
||||
rows, err := s.assets.pool.Query(ctx, `SELECT resource_type,resource_id::text FROM gateway.marketplace_installations WHERE portal_user_id=$1 ORDER BY created_at DESC`, portalUserID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
type bound struct{ typ, id string }
|
||||
bounds := []bound{}
|
||||
for rows.Next() {
|
||||
var b bound
|
||||
if err = rows.Scan(&b.typ, &b.id); err != nil {
|
||||
rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
bounds = append(bounds, b)
|
||||
}
|
||||
rows.Close()
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items := make([]MarketItem, 0, len(bounds))
|
||||
for _, b := range bounds {
|
||||
item, ok, err := s.resourceByID(ctx, b.typ, b.id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if ok {
|
||||
items = append(items, item)
|
||||
}
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// Installed reports whether a portal user has an installation for the resource.
|
||||
func (s *MarketplaceService) Installed(ctx context.Context, resourceType, resourceID, portalUserID string) (bool, error) {
|
||||
var exists bool
|
||||
err := s.assets.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.marketplace_installations WHERE resource_type=$1 AND resource_id=$2 AND portal_user_id=$3)`, resourceType, resourceID, portalUserID).Scan(&exists)
|
||||
return exists, err
|
||||
}
|
||||
|
||||
// PortalUserForAPIKey resolves the portal user that owns an API key, used to
|
||||
// evaluate installation-based visibility at runtime.
|
||||
func (s *MarketplaceService) PortalUserForAPIKey(ctx context.Context, apiKeyID string) (string, bool, error) {
|
||||
var userID *string
|
||||
err := s.assets.pool.QueryRow(ctx, `SELECT portal_user_id::text FROM gateway.api_keys WHERE id=$1`, apiKeyID).Scan(&userID)
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return "", false, nil
|
||||
}
|
||||
return "", false, err
|
||||
}
|
||||
if userID == nil || *userID == "" {
|
||||
return "", false, nil
|
||||
}
|
||||
return *userID, true, nil
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) publishedResourceID(ctx context.Context, resourceType, code string) (string, error) {
|
||||
var id string
|
||||
var err error
|
||||
switch resourceType {
|
||||
case "mcp_server":
|
||||
err = s.assets.pool.QueryRow(ctx, `SELECT id::text FROM gateway.mcp_servers WHERE code=$1 AND status='published' AND enabled`, code).Scan(&id)
|
||||
case "skill":
|
||||
err = s.assets.pool.QueryRow(ctx, `SELECT id::text FROM gateway.skills WHERE code=$1 AND status='published' AND enabled`, code).Scan(&id)
|
||||
case "digital_employee":
|
||||
err = s.assets.pool.QueryRow(ctx, `SELECT id::text FROM gateway.digital_employees WHERE code=$1 AND status='published' AND enabled`, code).Scan(&id)
|
||||
default:
|
||||
return "", errors.New("未知的资源类型")
|
||||
}
|
||||
return id, mapNotFound(err)
|
||||
}
|
||||
|
||||
func (s *MarketplaceService) resourceByID(ctx context.Context, resourceType, id string) (MarketItem, bool, error) {
|
||||
var item MarketItem
|
||||
switch resourceType {
|
||||
case "mcp_server":
|
||||
server, err := s.mcpServers.Get(ctx, id)
|
||||
if err != nil {
|
||||
return item, false, nil
|
||||
}
|
||||
item = MarketItem{Type: "mcp_server", Code: server.Code, Name: server.Name, Description: server.Description, CategoryID: server.CategoryID, CategoryName: server.CategoryName, Tags: server.Tags, UpdatedAt: server.UpdatedAt}
|
||||
case "skill":
|
||||
skill, err := s.skills.Get(ctx, id)
|
||||
if err != nil {
|
||||
return item, false, nil
|
||||
}
|
||||
item = MarketItem{Type: "skill", Code: skill.Code, Name: skill.Name, Description: skill.Description, CategoryID: skill.CategoryID, CategoryName: skill.CategoryName, Tags: skill.Tags, UpdatedAt: skill.UpdatedAt}
|
||||
case "digital_employee":
|
||||
employee, err := s.employees.Get(ctx, id)
|
||||
if err != nil {
|
||||
return item, false, nil
|
||||
}
|
||||
item = MarketItem{Type: "digital_employee", Code: employee.Code, Name: employee.Name, Description: employee.Description, CategoryID: employee.CategoryID, CategoryName: employee.CategoryName, Tags: employee.Tags, UpdatedAt: employee.UpdatedAt}
|
||||
default:
|
||||
return item, false, nil
|
||||
}
|
||||
return item, true, nil
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
"aigateway.local/core/internal/platform/database"
|
||||
)
|
||||
|
||||
func TestMarketplaceLifecycle(t *testing.T) {
|
||||
databaseURL := os.Getenv("WORKBENCH_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("WORKBENCH_TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8, MinConns: 0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
actorID := "22222222-2222-4222-8222-222222222222"
|
||||
portalUserID := "33333333-3333-4333-8333-333333333333"
|
||||
cleanup := func() {
|
||||
// pgx v5 Exec uses the extended protocol, which rejects multi-command
|
||||
// strings, so each statement must run in its own Exec.
|
||||
for _, q := range []string{
|
||||
`DELETE FROM gateway.marketplace_installations WHERE portal_user_id='` + portalUserID + `'`,
|
||||
`DELETE FROM gateway.digital_employees WHERE code='mkt_employee'`,
|
||||
`DELETE FROM gateway.skills WHERE code='mkt_skill'`,
|
||||
`DELETE FROM gateway.mcp_servers WHERE code='mkt_mcp'`,
|
||||
`DELETE FROM gateway.marketplace_categories WHERE name IN ('mkt-分类','mkt-技能分类')`,
|
||||
`DELETE FROM gateway.portal_users WHERE id='` + portalUserID + `'`,
|
||||
} {
|
||||
if _, err := pool.Exec(ctx, q); err != nil {
|
||||
t.Fatalf("cleanup: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
cleanup()
|
||||
defer cleanup()
|
||||
// Inserts must come after the initial cleanup, which would otherwise
|
||||
// delete the portal user we just created (it is referenced by the
|
||||
// marketplace_installations FK).
|
||||
if _, err := pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role) VALUES($1,'market-test','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `INSERT INTO gateway.portal_users(id,account,name) VALUES($1,'market-user','市场测试') ON CONFLICT(id) DO NOTHING`, portalUserID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
assets := NewService(pool)
|
||||
key := base64.StdEncoding.EncodeToString(make([]byte, 32))
|
||||
cipher, err := cryptox.NewKeyring(key, 1, "", "mkt-test-headers")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
mcpSvc := NewMCPServerService(assets, cipher, true)
|
||||
toolSvc := NewToolService(assets, cipher, true)
|
||||
skillSvc := NewSkillService(assets)
|
||||
deSvc := NewDigitalEmployeeService(assets, skillSvc, toolSvc, mcpSvc)
|
||||
market := NewMarketplaceService(assets, mcpSvc, skillSvc, deSvc)
|
||||
|
||||
globalCat, err := market.CreateCategory(ctx, Category{Name: "mkt-分类", Description: "全局", ResourceType: ""}, actorID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
skillCat, err := market.CreateCategory(ctx, Category{Name: "mkt-技能分类", Description: "技能专用", ResourceType: "skill"}, actorID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// A skill category must reject binding to an MCP server.
|
||||
if _, err := market.CreateCategory(ctx, Category{Name: "mkt-冲突分类", Description: "", ResourceType: "mcp_server"}, actorID); err != nil {
|
||||
t.Fatalf("mcp_server category should be creatable: %v", err)
|
||||
}
|
||||
cleanupConflict := func() {
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.marketplace_categories WHERE name='mkt-冲突分类'`)
|
||||
}
|
||||
cleanupConflict()
|
||||
defer cleanupConflict()
|
||||
|
||||
mcp, err := mcpSvc.Save(ctx, "", MCPServerInput{
|
||||
Code: "mkt_mcp", Name: "测试 MCP", Transport: "streamable-http",
|
||||
EndpointURL: "https://mcp.example.com/mcp", Status: "published",
|
||||
CategoryID: &globalCat.ID, Enabled: true, Headers: map[string]string{"Authorization": "Bearer test"},
|
||||
}, actorID, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
skill, err := skillSvc.Save(ctx, "", SkillInput{
|
||||
Code: "mkt_skill", Name: "测试技能", Content: "你是客服,请回答 {{question}}",
|
||||
Variables: []Variable{{Name: "question", Required: true}}, Status: "published",
|
||||
CategoryID: &skillCat.ID, Enabled: true, MCPServerIDs: []string{mcp.ID},
|
||||
}, actorID, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rendered, err := skillSvc.Render(skill, map[string]any{"question": "你好"})
|
||||
if err != nil || rendered != "你是客服,请回答 你好" {
|
||||
t.Fatalf("render=%q err=%v", rendered, err)
|
||||
}
|
||||
|
||||
employee, err := deSvc.Save(ctx, "", DigitalEmployeeInput{
|
||||
Code: "mkt_employee", Name: "客服机器人", Persona: "你是一位耐心的客服", Model: "deepseek-chat",
|
||||
SkillIDs: []string{skill.ID}, Status: "draft", Enabled: true,
|
||||
}, actorID, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if employee.Status != "draft" {
|
||||
t.Fatalf("expected draft, got %s", employee.Status)
|
||||
}
|
||||
published, err := deSvc.Publish(ctx, employee.ID, actorID)
|
||||
if err != nil {
|
||||
t.Fatalf("publish failed: %v", err)
|
||||
}
|
||||
if published.Status != "published" {
|
||||
t.Fatalf("expected published, got %s", published.Status)
|
||||
}
|
||||
|
||||
// Unified catalog merges all three published resource types.
|
||||
items, err := market.Catalog(ctx, "", "", "", "", 50)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := map[string]bool{}
|
||||
for _, item := range items {
|
||||
found[item.Code] = true
|
||||
}
|
||||
for _, code := range []string{"mkt_mcp", "mkt_skill", "mkt_employee"} {
|
||||
if !found[code] {
|
||||
t.Fatalf("catalog missing %s (got %#v)", code, found)
|
||||
}
|
||||
}
|
||||
// Category filter narrows to the skill category.
|
||||
skillItems, err := market.Catalog(ctx, "skill", skillCat.ID, "", "", 50)
|
||||
if err != nil || len(skillItems) != 1 || skillItems[0].Code != "mkt_skill" {
|
||||
t.Fatalf("category filter failed: %#v err=%v", skillItems, err)
|
||||
}
|
||||
|
||||
detailItem, raw, err := market.Detail(ctx, "digital_employee", "mkt_employee")
|
||||
if err != nil || detailItem.Code != "mkt_employee" || len(raw) == 0 {
|
||||
t.Fatalf("detail=%#v err=%v", detailItem, err)
|
||||
}
|
||||
|
||||
// Install/uninstall is idempotent and gated by published status.
|
||||
created, err := market.Install(ctx, "skill", "mkt_skill", portalUserID)
|
||||
if err != nil || !created {
|
||||
t.Fatalf("install created=%v err=%v", created, err)
|
||||
}
|
||||
installed, err := market.Installed(ctx, "skill", skill.ID, portalUserID)
|
||||
if err != nil || !installed {
|
||||
t.Fatalf("installed=%v err=%v", installed, err)
|
||||
}
|
||||
createdAgain, err := market.Install(ctx, "skill", "mkt_skill", portalUserID)
|
||||
if err != nil || createdAgain {
|
||||
t.Fatalf("re-install should be a no-op: created=%v err=%v", createdAgain, err)
|
||||
}
|
||||
installedList, err := market.ListInstalled(ctx, portalUserID)
|
||||
if err != nil || len(installedList) != 1 || installedList[0].Code != "mkt_skill" {
|
||||
t.Fatalf("installed list=%#v err=%v", installedList, err)
|
||||
}
|
||||
if err := market.Uninstall(ctx, "skill", "mkt_skill", portalUserID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
installed, err = market.Installed(ctx, "skill", skill.ID, portalUserID)
|
||||
if err != nil || installed {
|
||||
t.Fatalf("expected uninstalled, got installed=%v err=%v", installed, err)
|
||||
}
|
||||
|
||||
// Runs mirroring: publishing a digital employee whose skill is bound to an
|
||||
// MCP server must still revalidate cleanly (strict bindings).
|
||||
if _, err := deSvc.Publish(ctx, employee.ID, actorID); err != nil {
|
||||
t.Fatalf("re-publish after install should pass: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
"aigateway.local/core/internal/provider"
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type MCPServer struct {
|
||||
ID string `json:"id"`
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Transport string `json:"transport"`
|
||||
EndpointURL string `json:"endpoint_url"`
|
||||
Status string `json:"status"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
CategoryName string `json:"category_name"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
HasSecretHeaders bool `json:"has_secret_headers"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
EncryptedHeaders []byte `json:"-"`
|
||||
HeadersKEKVersion int `json:"-"`
|
||||
}
|
||||
|
||||
type MCPServerInput struct {
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Transport string `json:"transport"`
|
||||
EndpointURL string `json:"endpoint_url"`
|
||||
Headers map[string]string `json:"headers,omitempty"`
|
||||
Status string `json:"status"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
// MCPServerService manages registered MCP servers: lifecycle CRUD plus header
|
||||
// decryption for the MCP client.
|
||||
type MCPServerService struct {
|
||||
assets *Service
|
||||
cipher cryptox.Cipher
|
||||
allowPrivate bool
|
||||
}
|
||||
|
||||
func NewMCPServerService(assets *Service, cipher cryptox.Cipher, allowPrivate bool) *MCPServerService {
|
||||
return &MCPServerService{assets: assets, cipher: cipher, allowPrivate: allowPrivate}
|
||||
}
|
||||
|
||||
func (s *MCPServerService) validate(ctx context.Context, input *MCPServerInput, create bool) error {
|
||||
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Description = strings.TrimSpace(input.Description)
|
||||
input.Transport = strings.ToLower(strings.TrimSpace(input.Transport))
|
||||
if !codeRE.MatchString(input.Code) || input.Name == "" || len(input.Name) > 128 || len(input.Description) > 4000 {
|
||||
return errors.New("MCP 服务器编码、名称或描述格式无效")
|
||||
}
|
||||
switch input.Transport {
|
||||
case "", "streamable-http":
|
||||
input.Transport = "streamable-http"
|
||||
case "sse":
|
||||
default:
|
||||
return errors.New("不支持的 MCP 传输方式")
|
||||
}
|
||||
validated, err := provider.ValidateBaseURL(ctx, input.EndpointURL, s.allowPrivate)
|
||||
if err != nil {
|
||||
return fmt.Errorf("MCP 端点校验失败: %w", err)
|
||||
}
|
||||
input.EndpointURL = validated
|
||||
switch input.Status {
|
||||
case "", "draft":
|
||||
input.Status = "draft"
|
||||
case "published", "archived":
|
||||
default:
|
||||
return errors.New("无效的资源状态")
|
||||
}
|
||||
if input.CategoryID != nil && strings.TrimSpace(*input.CategoryID) == "" {
|
||||
input.CategoryID = nil
|
||||
}
|
||||
var categoryErr error
|
||||
input.CategoryID, categoryErr = validCategoryID(ctx, s.assets.pool, input.CategoryID, "")
|
||||
if categoryErr != nil {
|
||||
return categoryErr
|
||||
}
|
||||
input.Tags, err = normalizeStrings(input.Tags, 30)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
input.DepartmentIDs, err = normalizeStrings(input.DepartmentIDs, 100)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
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("MCP 请求头格式无效")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const mcpServerSelect = `SELECT m.id::text,m.code,m.name,m.description,m.transport,m.endpoint_url,m.status,m.category_id::text,coalesce(c.name,''),m.tags,m.department_ids::text[],m.enabled,octet_length(m.encrypted_headers)>0,m.revision,m.created_at,m.updated_at,m.encrypted_headers,m.headers_kek_version FROM gateway.mcp_servers m LEFT JOIN gateway.marketplace_categories c ON c.id=m.category_id`
|
||||
|
||||
func scanMCPServer(row pgx.Row) (MCPServer, error) {
|
||||
var s MCPServer
|
||||
err := row.Scan(&s.ID, &s.Code, &s.Name, &s.Description, &s.Transport, &s.EndpointURL, &s.Status, &s.CategoryID, &s.CategoryName, &s.Tags, &s.DepartmentIDs, &s.Enabled, &s.HasSecretHeaders, &s.Revision, &s.CreatedAt, &s.UpdatedAt, &s.EncryptedHeaders, &s.HeadersKEKVersion)
|
||||
return s, mapNotFound(err)
|
||||
}
|
||||
|
||||
func (s *MCPServerService) List(ctx context.Context) ([]MCPServer, error) {
|
||||
rows, err := s.assets.pool.Query(ctx, mcpServerSelect+` ORDER BY m.updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []MCPServer{}
|
||||
for rows.Next() {
|
||||
server, err := scanMCPServer(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, server)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *MCPServerService) Get(ctx context.Context, id string) (MCPServer, error) {
|
||||
return scanMCPServer(s.assets.pool.QueryRow(ctx, mcpServerSelect+` WHERE m.id=$1`, id))
|
||||
}
|
||||
|
||||
func (s *MCPServerService) GetByCode(ctx context.Context, code string) (MCPServer, error) {
|
||||
return scanMCPServer(s.assets.pool.QueryRow(ctx, mcpServerSelect+` WHERE m.code=$1`, code))
|
||||
}
|
||||
|
||||
// GetPublishedByCode returns a published, enabled server by code.
|
||||
func (s *MCPServerService) GetPublishedByCode(ctx context.Context, code string) (MCPServer, error) {
|
||||
return scanMCPServer(s.assets.pool.QueryRow(ctx, mcpServerSelect+` WHERE m.code=$1 AND m.status='published' AND m.enabled`, code))
|
||||
}
|
||||
|
||||
func (s *MCPServerService) Save(ctx context.Context, id string, input MCPServerInput, actorID string, create bool) (MCPServer, error) {
|
||||
if err := s.validate(ctx, &input, create); err != nil {
|
||||
return MCPServer{}, err
|
||||
}
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return MCPServer{}, 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 MCPServer{}, err
|
||||
}
|
||||
}
|
||||
if create {
|
||||
id, err = newUUID()
|
||||
if err != nil {
|
||||
return MCPServer{}, err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.mcp_servers(id,code,name,description,transport,endpoint_url,encrypted_headers,headers_kek_version,status,category_id,tags,department_ids,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)`, id, input.Code, input.Name, input.Description, input.Transport, input.EndpointURL, encrypted, version, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled, actorID)
|
||||
} else if input.Headers == nil {
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.mcp_servers SET code=$2,name=$3,description=$4,transport=$5,endpoint_url=$6,status=$7,category_id=$8,tags=$9,department_ids=$10,enabled=$11,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.Transport, input.EndpointURL, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return MCPServer{}, ErrNotFound
|
||||
}
|
||||
} else {
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.mcp_servers SET code=$2,name=$3,description=$4,transport=$5,endpoint_url=$6,encrypted_headers=$7,headers_kek_version=$8,status=$9,category_id=$10,tags=$11,department_ids=$12,enabled=$13,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.Transport, input.EndpointURL, encrypted, version, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return MCPServer{}, ErrNotFound
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return MCPServer{}, err
|
||||
}
|
||||
event := "mcp_server.updated"
|
||||
if create {
|
||||
event = "mcp_server.created"
|
||||
}
|
||||
if err = emit(ctx, tx, event, "mcp_server", id, actorID, nil); err != nil {
|
||||
return MCPServer{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return MCPServer{}, err
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
func (s *MCPServerService) 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.skills WHERE $1 = ANY(mcp_server_ids)
|
||||
UNION ALL SELECT 1 FROM gateway.digital_employees WHERE $1 = ANY(mcp_server_ids)
|
||||
UNION ALL SELECT 1 FROM gateway.marketplace_installations WHERE resource_type='mcp_server' AND resource_id=$1
|
||||
)`, id).Scan(&used); err != nil {
|
||||
return err
|
||||
}
|
||||
if used {
|
||||
return ErrConflict
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.mcp_servers WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "mcp_server.deleted", "mcp_server", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
// Headers decrypts the stored request headers for an MCP server so the client
|
||||
// can authenticate to it.
|
||||
func (s *MCPServerService) Headers(server MCPServer) (map[string]string, error) {
|
||||
if len(server.EncryptedHeaders) == 0 {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
plain, err := s.cipher.Decrypt(server.EncryptedHeaders, server.HeadersKEKVersion)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
headers := map[string]string{}
|
||||
if err = json.Unmarshal(plain, &headers); err != nil {
|
||||
return nil, errors.New("MCP 请求头密文内容无效")
|
||||
}
|
||||
return headers, nil
|
||||
}
|
||||
@@ -0,0 +1,369 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// mcpProtocolVersion is the Model Context Protocol version this client speaks.
|
||||
// Streamable HTTP (2025-06-18) is the current stable transport; the server may
|
||||
// negotiate an older one and we accept whatever it replies with.
|
||||
const mcpProtocolVersion = "2025-06-18"
|
||||
|
||||
// MCPTool is one tool discovered from an MCP server via tools/list.
|
||||
type MCPTool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
InputSchema json.RawMessage `json:"inputSchema"`
|
||||
}
|
||||
|
||||
// MCPToolResult is the outcome of a tools/call.
|
||||
type MCPToolResult struct {
|
||||
Content string
|
||||
IsError bool
|
||||
}
|
||||
|
||||
// mcpServerState caches the negotiated session and discovered tool list for one
|
||||
// server, keyed by server ID. initialize is expensive (two round trips) and
|
||||
// stateless HTTP servers expect it per session, so we keep the session warm for
|
||||
// cacheTTL and only re-handshake once it lapses.
|
||||
type mcpServerState struct {
|
||||
initAt time.Time
|
||||
sessionID string
|
||||
tools []MCPTool
|
||||
toolsAt time.Time
|
||||
}
|
||||
|
||||
// MCPClient is a minimal Model Context Protocol client over the streamable HTTP
|
||||
// (and legacy SSE) transports. It speaks just enough of the protocol for the
|
||||
// resource marketplace: initialize + notifications/initialized, tools/list for
|
||||
// discovery, and tools/call for execution.
|
||||
type MCPClient struct {
|
||||
client *http.Client
|
||||
cacheTTL time.Duration
|
||||
mu sync.Mutex
|
||||
states map[string]*mcpServerState
|
||||
}
|
||||
|
||||
func NewMCPClient(allowPrivate bool, cacheTTL time.Duration) *MCPClient {
|
||||
if cacheTTL <= 0 {
|
||||
cacheTTL = 60 * time.Second
|
||||
}
|
||||
return &MCPClient{
|
||||
client: &http.Client{
|
||||
Timeout: 15 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
DialContext: safeToolDial(allowPrivate),
|
||||
ForceAttemptHTTP2: true,
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
ResponseHeaderTimeout: 10 * time.Second,
|
||||
MaxIdleConns: 64,
|
||||
MaxIdleConnsPerHost: 16,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
},
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("MCP 服务器不允许重定向") },
|
||||
},
|
||||
cacheTTL: cacheTTL,
|
||||
states: make(map[string]*mcpServerState),
|
||||
}
|
||||
}
|
||||
|
||||
// DiscoverTools returns the tools an MCP server advertises, reusing a cached
|
||||
// list for cacheTTL. headers are the already-decrypted request headers (e.g.
|
||||
// Authorization) for this server.
|
||||
func (c *MCPClient) DiscoverTools(ctx context.Context, server MCPServer, headers map[string]string) ([]MCPTool, error) {
|
||||
state, err := c.ensureInitialized(ctx, server, headers)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if state.tools != nil && time.Since(state.toolsAt) < c.cacheTTL {
|
||||
return state.tools, nil
|
||||
}
|
||||
result, err := c.call(ctx, server, headers, "tools/list", map[string]any{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var list struct {
|
||||
Tools []MCPTool `json:"tools"`
|
||||
}
|
||||
if err = json.Unmarshal(result, &list); err != nil {
|
||||
return nil, fmt.Errorf("MCP tools/list 响应无效: %w", err)
|
||||
}
|
||||
for i := range list.Tools {
|
||||
if list.Tools[i].Name == "" {
|
||||
return nil, errors.New("MCP 服务器返回了没有名称的工具")
|
||||
}
|
||||
if len(list.Tools[i].InputSchema) == 0 {
|
||||
list.Tools[i].InputSchema = json.RawMessage(`{}`)
|
||||
}
|
||||
}
|
||||
c.mu.Lock()
|
||||
state.tools = list.Tools
|
||||
state.toolsAt = time.Now()
|
||||
c.mu.Unlock()
|
||||
return list.Tools, nil
|
||||
}
|
||||
|
||||
// CallTool invokes one tool on an MCP server and returns the concatenated text
|
||||
// content. An isError result is surfaced as an error so callers treat it as a
|
||||
// failed tool round rather than a successful empty answer.
|
||||
func (c *MCPClient) CallTool(ctx context.Context, server MCPServer, headers map[string]string, name string, args map[string]any) (MCPToolResult, error) {
|
||||
state, err := c.ensureInitialized(ctx, server, headers)
|
||||
if err != nil {
|
||||
return MCPToolResult{}, err
|
||||
}
|
||||
// The runtime exposes tools under the collision-proof prefix
|
||||
// (mcp__{serverCode}__{toolName}); strip it before the wire call since the
|
||||
// remote server only knows the unprefixed tool name.
|
||||
if _, resolved, ok := resolveMCPTool(name); ok {
|
||||
name = resolved
|
||||
}
|
||||
params := map[string]any{"name": name, "arguments": args}
|
||||
result, err := c.call(ctx, server, headers, "tools/call", params)
|
||||
if err != nil {
|
||||
return MCPToolResult{}, err
|
||||
}
|
||||
var called struct {
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
IsError bool `json:"isError"`
|
||||
}
|
||||
if err = json.Unmarshal(result, &called); err != nil {
|
||||
return MCPToolResult{}, fmt.Errorf("MCP tools/call 响应无效: %w", err)
|
||||
}
|
||||
text := strings.Builder{}
|
||||
for _, block := range called.Content {
|
||||
if block.Type == "text" {
|
||||
if text.Len() > 0 {
|
||||
text.WriteString("\n")
|
||||
}
|
||||
text.WriteString(block.Text)
|
||||
}
|
||||
}
|
||||
_ = state // keep state warm check semantics obvious
|
||||
if called.IsError {
|
||||
return MCPToolResult{}, errors.New("MCP 工具执行失败: " + text.String())
|
||||
}
|
||||
return MCPToolResult{Content: text.String()}, nil
|
||||
}
|
||||
|
||||
// ensureInitialized performs the MCP initialize handshake for a server if its
|
||||
// session has lapsed (or no cached tools exist yet), then acknowledges with
|
||||
// notifications/initialized. The handshake is guarded by the per-server cache
|
||||
// so a burst of calls does not re-initialize every request.
|
||||
func (c *MCPClient) ensureInitialized(ctx context.Context, server MCPServer, headers map[string]string) (*mcpServerState, error) {
|
||||
c.mu.Lock()
|
||||
state, ok := c.states[server.ID]
|
||||
if ok && time.Since(state.initAt) < c.cacheTTL {
|
||||
c.mu.Unlock()
|
||||
return state, nil
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
result, headersOut, err := c.handshake(ctx, server, headers)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sessionID := headersOut.Get("Mcp-Session-Id")
|
||||
if server.Transport != "" && server.Transport != "streamable-http" && server.Transport != "sse" {
|
||||
return nil, fmt.Errorf("不支持的 MCP 传输方式 %s", server.Transport)
|
||||
}
|
||||
_ = result // negotiated protocol version is accepted as-is
|
||||
|
||||
c.mu.Lock()
|
||||
state = &mcpServerState{initAt: time.Now(), sessionID: sessionID}
|
||||
c.states[server.ID] = state
|
||||
c.mu.Unlock()
|
||||
|
||||
// Best-effort acknowledgment; servers that require it will reject later
|
||||
// calls and we will surface that error naturally.
|
||||
c.sendNotification(ctx, server, headers, sessionID)
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) handshake(ctx context.Context, server MCPServer, headers map[string]string) (json.RawMessage, http.Header, error) {
|
||||
params := map[string]any{
|
||||
"protocolVersion": mcpProtocolVersion,
|
||||
"capabilities": map[string]any{},
|
||||
"clientInfo": map[string]any{"name": "LLMGuardX语枢", "version": "0.10.0"},
|
||||
}
|
||||
payload, _ := json.Marshal(mcpRequest{JSONRPC: "2.0", ID: 1, Method: "initialize", Params: params})
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, server.EndpointURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
c.prepare(request, headers, "")
|
||||
response, err := c.client.Do(request)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("MCP 连接失败: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
result, err := readMCPBody(response)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return result, response.Header, nil
|
||||
}
|
||||
|
||||
func (c *MCPClient) sendNotification(ctx context.Context, server MCPServer, headers map[string]string, sessionID string) {
|
||||
payload, _ := json.Marshal(mcpRequest{JSONRPC: "2.0", ID: nil, Method: "notifications/initialized"})
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, server.EndpointURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
c.prepare(request, headers, sessionID)
|
||||
response, err := c.client.Do(request)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 1<<10))
|
||||
response.Body.Close()
|
||||
}
|
||||
|
||||
func (c *MCPClient) call(ctx context.Context, server MCPServer, headers map[string]string, method string, params any) (json.RawMessage, error) {
|
||||
c.mu.Lock()
|
||||
sessionID := ""
|
||||
if state, ok := c.states[server.ID]; ok {
|
||||
sessionID = state.sessionID
|
||||
}
|
||||
c.mu.Unlock()
|
||||
payload, _ := json.Marshal(mcpRequest{JSONRPC: "2.0", ID: 1, Method: method, Params: params})
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, server.EndpointURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.prepare(request, headers, sessionID)
|
||||
response, err := c.client.Do(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("MCP 调用失败: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
return readMCPBody(response)
|
||||
}
|
||||
|
||||
func (c *MCPClient) prepare(request *http.Request, headers map[string]string, sessionID string) {
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("Accept", "application/json, text/event-stream")
|
||||
for key, value := range headers {
|
||||
request.Header.Set(key, value)
|
||||
}
|
||||
if sessionID != "" {
|
||||
request.Header.Set("Mcp-Session-Id", sessionID)
|
||||
}
|
||||
}
|
||||
|
||||
type mcpRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID any `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params any `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
type mcpError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type mcpResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID any `json:"id"`
|
||||
Result json.RawMessage `json:"result"`
|
||||
Error *mcpError `json:"error"`
|
||||
}
|
||||
|
||||
// readMCPBody reads and parses a single JSON-RPC response. Streamable HTTP
|
||||
// servers return application/json; legacy SSE servers stream data frames, from
|
||||
// which the first complete JSON object is extracted.
|
||||
func readMCPBody(response *http.Response) (json.RawMessage, error) {
|
||||
code := response.StatusCode
|
||||
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("MCP 响应超过 1 MiB")
|
||||
}
|
||||
if code < 200 || code >= 300 {
|
||||
return nil, fmt.Errorf("MCP 服务器返回 HTTP %d", code)
|
||||
}
|
||||
var body []byte
|
||||
if strings.HasPrefix(response.Header.Get("Content-Type"), "text/event-stream") {
|
||||
body, err = extractSSEJSON(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else {
|
||||
body = raw
|
||||
}
|
||||
var resp mcpResponse
|
||||
if err = json.Unmarshal(body, &resp); err != nil {
|
||||
return nil, fmt.Errorf("MCP 响应不是有效 JSON-RPC: %w", err)
|
||||
}
|
||||
if resp.Error != nil {
|
||||
return nil, fmt.Errorf("MCP 服务器错误 (%d): %s", resp.Error.Code, resp.Error.Message)
|
||||
}
|
||||
if len(resp.Result) == 0 {
|
||||
return nil, errors.New("MCP 服务器返回空结果")
|
||||
}
|
||||
return resp.Result, nil
|
||||
}
|
||||
|
||||
// extractSSEJSON concatenates the first data frame's payload into one JSON
|
||||
// document. MCP streamable-HTTP servers emit a single frame per request, so
|
||||
// only the first frame boundary is consumed and the rest is ignored.
|
||||
func extractSSEJSON(raw []byte) ([]byte, error) {
|
||||
var data strings.Builder
|
||||
for _, line := range strings.Split(string(raw), "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
switch {
|
||||
case strings.HasPrefix(trimmed, "data:"):
|
||||
value := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
|
||||
if value == "" {
|
||||
continue
|
||||
}
|
||||
if data.Len() > 0 {
|
||||
data.WriteString("\n")
|
||||
}
|
||||
data.WriteString(value)
|
||||
case trimmed == "" && data.Len() > 0:
|
||||
if json.Valid([]byte(data.String())) {
|
||||
return []byte(data.String()), nil
|
||||
}
|
||||
data.Reset()
|
||||
}
|
||||
}
|
||||
if data.Len() > 0 && json.Valid([]byte(data.String())) {
|
||||
return []byte(data.String()), nil
|
||||
}
|
||||
return nil, errors.New("MCP 服务器未返回有效的 SSE 数据帧")
|
||||
}
|
||||
|
||||
// toolCallPrefix namespaces MCP tools inside a shared tool list so different
|
||||
// servers cannot collide. Format: mcp__{serverCode}__{toolName}.
|
||||
const toolCallPrefix = "mcp__"
|
||||
|
||||
func mcpToolName(serverCode, tool string) string { return toolCallPrefix + serverCode + "__" + tool }
|
||||
|
||||
// resolveMCPTool splits a prefixed tool name back into its MCP server code and
|
||||
// tool name. Returns ok=false for names that are not MCP-prefixed.
|
||||
func resolveMCPTool(prefixed string) (serverCode, tool string, ok bool) {
|
||||
if !strings.HasPrefix(prefixed, toolCallPrefix) {
|
||||
return "", "", false
|
||||
}
|
||||
rest := strings.TrimPrefix(prefixed, toolCallPrefix)
|
||||
parts := strings.SplitN(rest, "__", 2)
|
||||
if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
|
||||
return "", "", false
|
||||
}
|
||||
return parts[0], parts[1], true
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// minimalMCP is a scriptable JSON-RPC MCP server for exercising the client.
|
||||
type minimalMCP struct {
|
||||
t *testing.T
|
||||
server *httptest.Server
|
||||
tools []MCPTool
|
||||
initCount int32
|
||||
listCount int32
|
||||
callCount int32
|
||||
lastArgs map[string]any
|
||||
lastTool string
|
||||
versioned bool
|
||||
}
|
||||
|
||||
func newMinimalMCP(t *testing.T) *minimalMCP {
|
||||
m := &minimalMCP{t: t, tools: []MCPTool{
|
||||
{Name: "lookup", Description: "look something up", InputSchema: json.RawMessage(`{"type":"object","properties":{"q":{"type":"string"}},"required":["q"]}`)},
|
||||
{Name: "add", Description: "add two numbers", InputSchema: json.RawMessage(`{"type":"object","properties":{"a":{"type":"number"},"b":{"type":"number"}},"required":["a","b"]}`)},
|
||||
}}
|
||||
m.server = httptest.NewServer(http.HandlerFunc(m.serve))
|
||||
t.Cleanup(m.server.Close)
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *minimalMCP) endpoint() string { return m.server.URL + "/mcp" }
|
||||
|
||||
func (m *minimalMCP) serve(w http.ResponseWriter, r *http.Request) {
|
||||
var req mcpRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "bad json", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
atomic.AddInt32(&m.initCount, 1)
|
||||
params, _ := req.Params.(map[string]any)
|
||||
clientInfo, _ := params["clientInfo"].(map[string]any)
|
||||
if clientInfo["name"] == nil || clientInfo["name"] == "" {
|
||||
m.t.Error("initialize missing clientInfo.name")
|
||||
}
|
||||
result := map[string]any{
|
||||
"protocolVersion": "2025-06-18",
|
||||
"capabilities": map[string]any{"tools": map[string]any{}},
|
||||
"serverInfo": map[string]any{"name": "test-mcp", "version": "1.0.0"},
|
||||
}
|
||||
m.write(w, req.ID, result, false)
|
||||
case "notifications/initialized":
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
case "tools/list":
|
||||
atomic.AddInt32(&m.listCount, 1)
|
||||
m.write(w, req.ID, map[string]any{"tools": m.tools}, false)
|
||||
case "tools/call":
|
||||
atomic.AddInt32(&m.callCount, 1)
|
||||
params, _ := req.Params.(map[string]any)
|
||||
name, _ := params["name"].(string)
|
||||
args, _ := params["arguments"].(map[string]any)
|
||||
m.lastArgs = args
|
||||
m.lastTool = name
|
||||
if name == "boom" {
|
||||
m.write(w, req.ID, map[string]any{"content": []map[string]any{{"type": "text", "text": "failed intentionally"}}, "isError": true}, false)
|
||||
return
|
||||
}
|
||||
if name == "nope" {
|
||||
m.write(w, req.ID, nil, true)
|
||||
return
|
||||
}
|
||||
m.write(w, req.ID, map[string]any{"content": []map[string]any{{"type": "text", "text": "result for " + name}}, "isError": false}, false)
|
||||
case "tools/fail":
|
||||
m.write(w, req.ID, nil, true)
|
||||
default:
|
||||
http.Error(w, "unknown method", http.StatusBadRequest)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *minimalMCP) write(w http.ResponseWriter, id any, result any, isError bool) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if isError {
|
||||
json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, "error": map[string]any{"code": -32601, "message": "method not found"}})
|
||||
return
|
||||
}
|
||||
json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, "result": result})
|
||||
}
|
||||
|
||||
func testServer(code string) MCPServer {
|
||||
return MCPServer{ID: "00000000-0000-0000-0000-000000000001", Code: code}
|
||||
}
|
||||
|
||||
func TestMCPDiscoverTools(t *testing.T) {
|
||||
mcp := newMinimalMCP(t)
|
||||
client := NewMCPClient(true, time.Minute)
|
||||
server := testServer("demo")
|
||||
server.EndpointURL = mcp.endpoint()
|
||||
|
||||
tools, err := client.DiscoverTools(context.Background(), server, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("discover failed: %v", err)
|
||||
}
|
||||
if len(tools) != 2 || tools[0].Name != "lookup" {
|
||||
t.Fatalf("expected 2 tools starting with lookup, got %+v", tools)
|
||||
}
|
||||
if atomic.LoadInt32(&mcp.initCount) != 1 {
|
||||
t.Fatalf("expected exactly one initialize handshake, got %d", mcp.initCount)
|
||||
}
|
||||
// Second call hits the tool cache, no new tools/list.
|
||||
if _, err = client.DiscoverTools(context.Background(), server, nil); err != nil {
|
||||
t.Fatalf("cached discover failed: %v", err)
|
||||
}
|
||||
if atomic.LoadInt32(&mcp.listCount) != 1 {
|
||||
t.Fatalf("expected tools/list to run once, got %d", mcp.listCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPCacheExpiryRediscover(t *testing.T) {
|
||||
mcp := newMinimalMCP(t)
|
||||
client := NewMCPClient(true, 50*time.Millisecond)
|
||||
server := testServer("expire")
|
||||
server.EndpointURL = mcp.endpoint()
|
||||
|
||||
if _, err := client.DiscoverTools(context.Background(), server, nil); err != nil {
|
||||
t.Fatalf("discover failed: %v", err)
|
||||
}
|
||||
time.Sleep(80 * time.Millisecond)
|
||||
if _, err := client.DiscoverTools(context.Background(), server, nil); err != nil {
|
||||
t.Fatalf("rediscover failed: %v", err)
|
||||
}
|
||||
if atomic.LoadInt32(&mcp.listCount) != 2 {
|
||||
t.Fatalf("expected tools/list to run twice after expiry, got %d", mcp.listCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPCallTool(t *testing.T) {
|
||||
mcp := newMinimalMCP(t)
|
||||
client := NewMCPClient(true, time.Minute)
|
||||
server := testServer("call")
|
||||
server.EndpointURL = mcp.endpoint()
|
||||
|
||||
result, err := client.CallTool(context.Background(), server, nil, "lookup", map[string]any{"q": "hello"})
|
||||
if err != nil {
|
||||
t.Fatalf("call failed: %v", err)
|
||||
}
|
||||
if result.Content != "result for lookup" {
|
||||
t.Fatalf("unexpected content %q", result.Content)
|
||||
}
|
||||
if mcp.lastArgs["q"] != "hello" {
|
||||
t.Fatalf("arguments not forwarded: %v", mcp.lastArgs)
|
||||
}
|
||||
// The runtime exposes tools under the mcp__{code}__{tool} prefix; CallTool
|
||||
// must strip it so the remote server sees the real tool name.
|
||||
if mcp.lastTool != "lookup" {
|
||||
t.Fatalf("tool name not sent as-is: %q", mcp.lastTool)
|
||||
}
|
||||
if _, err := client.CallTool(context.Background(), server, nil, "mcp__call__lookup", map[string]any{"q": "x"}); err != nil {
|
||||
t.Fatalf("prefixed call failed: %v", err)
|
||||
}
|
||||
if mcp.lastTool != "lookup" {
|
||||
t.Fatalf("prefixed tool name was not stripped: %q", mcp.lastTool)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPCallToolIsError(t *testing.T) {
|
||||
mcp := newMinimalMCP(t)
|
||||
client := NewMCPClient(true, time.Minute)
|
||||
server := testServer("err")
|
||||
server.EndpointURL = mcp.endpoint()
|
||||
|
||||
if _, err := client.CallTool(context.Background(), server, nil, "boom", nil); err == nil {
|
||||
t.Fatal("expected error for isError tool result")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPJSONRPCError(t *testing.T) {
|
||||
mcp := newMinimalMCP(t)
|
||||
client := NewMCPClient(true, time.Minute)
|
||||
server := testServer("rpc")
|
||||
server.EndpointURL = mcp.endpoint()
|
||||
|
||||
if _, err := client.CallTool(context.Background(), server, nil, "nope", nil); err == nil || !strings.Contains(err.Error(), "method not found") {
|
||||
t.Fatalf("expected JSON-RPC error surfaced, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPToolNamePrefixRoundTrip(t *testing.T) {
|
||||
prefixed := mcpToolName("github", "create-issue")
|
||||
if prefixed != "mcp__github__create-issue" {
|
||||
t.Fatalf("unexpected prefix %q", prefixed)
|
||||
}
|
||||
code, tool, ok := resolveMCPTool(prefixed)
|
||||
if !ok || code != "github" || tool != "create-issue" {
|
||||
t.Fatalf("resolve failed: %q %q %v", code, tool, ok)
|
||||
}
|
||||
if _, _, ok := resolveMCPTool("plain-name"); ok {
|
||||
t.Fatal("non-prefixed name should not resolve as MCP tool")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPSSESingleFrame(t *testing.T) {
|
||||
handler := func(w http.ResponseWriter, r *http.Request) {
|
||||
var req mcpRequest
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
payload, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": map[string]any{"protocolVersion": "2025-06-18", "capabilities": map[string]any{}, "serverInfo": map[string]any{"name": "sse", "version": "1"}}})
|
||||
w.Write([]byte("event: message\ndata: " + string(payload) + "\n\n"))
|
||||
case "notifications/initialized":
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
case "tools/list":
|
||||
payload, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": map[string]any{"tools": []MCPTool{{Name: "sse-tool", Description: "", InputSchema: json.RawMessage(`{}`)}}}})
|
||||
w.Write([]byte("event: message\ndata: " + string(payload) + "\n\n"))
|
||||
}
|
||||
}
|
||||
server := httptest.NewServer(http.HandlerFunc(handler))
|
||||
defer server.Close()
|
||||
|
||||
client := NewMCPClient(true, time.Minute)
|
||||
svr := testServer("sse")
|
||||
svr.EndpointURL = server.URL
|
||||
tools, err := client.DiscoverTools(context.Background(), svr, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("SSE discover failed: %v", err)
|
||||
}
|
||||
if len(tools) != 1 || tools[0].Name != "sse-tool" {
|
||||
t.Fatalf("unexpected SSE tools: %+v", tools)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,437 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
"aigateway.local/core/internal/provider"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type NotificationInput struct {
|
||||
Name, WebhookURL string
|
||||
SigningSecret *string
|
||||
EventPatterns []string
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
type NotificationService struct {
|
||||
assets *Service
|
||||
cipher cryptox.Cipher
|
||||
allowPrivate bool
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func NewNotificationService(assets *Service, cipher cryptox.Cipher, allowPrivate bool) *NotificationService {
|
||||
// One shared client for all deliveries: http.Client is safe for concurrent
|
||||
// use, and reusing the transport keeps TLS sessions and connections
|
||||
// warm instead of re-handshaking for every webhook.
|
||||
client := &http.Client{
|
||||
Timeout: 15 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
DialContext: safeToolDial(allowPrivate),
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
ResponseHeaderTimeout: 10 * time.Second,
|
||||
MaxIdleConns: 64,
|
||||
MaxIdleConnsPerHost: 16,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
},
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("Webhook 不允许重定向") },
|
||||
}
|
||||
return &NotificationService{assets: assets, cipher: cipher, allowPrivate: allowPrivate, client: client}
|
||||
}
|
||||
|
||||
const channelSelect = `SELECT id::text,name,webhook_url,event_patterns,enabled,octet_length(encrypted_signing_secret)>0,revision,created_at,updated_at,encrypted_signing_secret,signing_secret_kek_version FROM gateway.notification_channels`
|
||||
|
||||
func scanChannel(row pgx.Row) (NotificationChannel, error) {
|
||||
var c NotificationChannel
|
||||
err := row.Scan(&c.ID, &c.Name, &c.WebhookURL, &c.EventPatterns, &c.Enabled, &c.HasSigningSecret, &c.Revision, &c.CreatedAt, &c.UpdatedAt, &c.EncryptedSigningSecret, &c.SigningSecretKEKVersion)
|
||||
return c, mapNotFound(err)
|
||||
}
|
||||
func (s *NotificationService) ListChannels(ctx context.Context) ([]NotificationChannel, error) {
|
||||
rows, err := s.assets.pool.Query(ctx, channelSelect+` ORDER BY updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []NotificationChannel{}
|
||||
for rows.Next() {
|
||||
c, err := scanChannel(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, c)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
func (s *NotificationService) GetChannel(ctx context.Context, id string) (NotificationChannel, error) {
|
||||
return scanChannel(s.assets.pool.QueryRow(ctx, channelSelect+` WHERE id=$1`, id))
|
||||
}
|
||||
|
||||
func (s *NotificationService) SaveChannel(ctx context.Context, id string, input NotificationInput, actorID string, create bool) (NotificationChannel, error) {
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.WebhookURL = strings.TrimSpace(input.WebhookURL)
|
||||
if input.Name == "" || len(input.Name) > 128 {
|
||||
return NotificationChannel{}, errors.New("通知通道名称格式无效")
|
||||
}
|
||||
validated, err := provider.ValidateBaseURL(ctx, input.WebhookURL, s.allowPrivate)
|
||||
if err != nil {
|
||||
return NotificationChannel{}, fmt.Errorf("Webhook 地址校验失败: %w", err)
|
||||
}
|
||||
input.WebhookURL = validated
|
||||
input.EventPatterns, err = normalizeStrings(input.EventPatterns, 100)
|
||||
if err != nil {
|
||||
return NotificationChannel{}, err
|
||||
}
|
||||
if len(input.EventPatterns) == 0 {
|
||||
return NotificationChannel{}, errors.New("至少配置一个事件模式")
|
||||
}
|
||||
for _, pattern := range input.EventPatterns {
|
||||
if strings.Count(pattern, "*") > 1 || (strings.Contains(pattern, "*") && !strings.HasSuffix(pattern, "*")) {
|
||||
return NotificationChannel{}, errors.New("事件模式仅允许末尾 * 通配")
|
||||
}
|
||||
}
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return NotificationChannel{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
var encrypted []byte
|
||||
var version int
|
||||
if input.SigningSecret != nil && *input.SigningSecret != "" {
|
||||
if len(*input.SigningSecret) > 4096 {
|
||||
return NotificationChannel{}, errors.New("签名密钥过长")
|
||||
}
|
||||
encrypted, version, err = s.cipher.Encrypt([]byte(*input.SigningSecret))
|
||||
if err != nil {
|
||||
return NotificationChannel{}, err
|
||||
}
|
||||
}
|
||||
if create {
|
||||
id, err = newUUID()
|
||||
if err != nil {
|
||||
return NotificationChannel{}, err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.notification_channels(id,name,webhook_url,encrypted_signing_secret,signing_secret_kek_version,event_patterns,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8)`, id, input.Name, input.WebhookURL, encrypted, version, input.EventPatterns, input.Enabled, actorID)
|
||||
} else if input.SigningSecret == nil {
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.notification_channels SET name=$2,webhook_url=$3,event_patterns=$4,enabled=$5,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Name, input.WebhookURL, input.EventPatterns, input.Enabled)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return NotificationChannel{}, ErrNotFound
|
||||
}
|
||||
} else {
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.notification_channels SET name=$2,webhook_url=$3,encrypted_signing_secret=$4,signing_secret_kek_version=$5,event_patterns=$6,enabled=$7,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Name, input.WebhookURL, encrypted, version, input.EventPatterns, input.Enabled)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return NotificationChannel{}, ErrNotFound
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return NotificationChannel{}, err
|
||||
}
|
||||
event := "notification_channel.updated"
|
||||
if create {
|
||||
event = "notification_channel.created"
|
||||
}
|
||||
if err = emit(ctx, tx, event, "notification_channel", id, actorID, nil); err != nil {
|
||||
return NotificationChannel{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return NotificationChannel{}, err
|
||||
}
|
||||
return s.GetChannel(ctx, id)
|
||||
}
|
||||
|
||||
func (s *NotificationService) DeleteChannel(ctx context.Context, id, actorID string) error {
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.notification_channels WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "notification_channel.deleted", "notification_channel", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
const deliverySelect = `SELECT d.id::text,d.channel_id::text,c.name,d.event_id::text,d.event_type,d.status,d.last_error,d.payload,d.attempts,d.response_status,d.delivered_at,d.created_at,d.updated_at FROM gateway.notification_deliveries d JOIN gateway.notification_channels c ON c.id=d.channel_id`
|
||||
|
||||
func (s *NotificationService) ListDeliveries(ctx context.Context, limit int) ([]NotificationDelivery, error) {
|
||||
if limit < 1 {
|
||||
limit = 100
|
||||
}
|
||||
if limit > 500 {
|
||||
limit = 500
|
||||
}
|
||||
rows, err := s.assets.pool.Query(ctx, deliverySelect+` ORDER BY d.created_at DESC LIMIT $1`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []NotificationDelivery{}
|
||||
for rows.Next() {
|
||||
var d NotificationDelivery
|
||||
if err = rows.Scan(&d.ID, &d.ChannelID, &d.ChannelName, &d.EventID, &d.EventType, &d.Status, &d.LastError, &d.Payload, &d.Attempts, &d.ResponseStatus, &d.DeliveredAt, &d.CreatedAt, &d.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, d)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func matchesEvent(patterns []string, eventType string) bool {
|
||||
for _, pattern := range patterns {
|
||||
if pattern == "*" || pattern == eventType {
|
||||
return true
|
||||
}
|
||||
if strings.HasSuffix(pattern, "*") && strings.HasPrefix(eventType, strings.TrimSuffix(pattern, "*")) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *NotificationService) ensureDelivery(ctx context.Context, channel NotificationChannel, eventID, eventType string, payload json.RawMessage) (NotificationDelivery, error) {
|
||||
id, err := newUUID()
|
||||
if err != nil {
|
||||
return NotificationDelivery{}, err
|
||||
}
|
||||
var d NotificationDelivery
|
||||
err = s.assets.pool.QueryRow(ctx, `INSERT INTO gateway.notification_deliveries(id,channel_id,event_id,event_type,payload,status) VALUES($1,$2,$3,$4,$5,'pending') ON CONFLICT(channel_id,event_id) DO UPDATE SET updated_at=gateway.notification_deliveries.updated_at RETURNING id::text,channel_id::text,$6,event_id::text,event_type,status,last_error,payload,attempts,response_status,delivered_at,created_at,updated_at`, id, channel.ID, eventID, eventType, payload, channel.Name).Scan(&d.ID, &d.ChannelID, &d.ChannelName, &d.EventID, &d.EventType, &d.Status, &d.LastError, &d.Payload, &d.Attempts, &d.ResponseStatus, &d.DeliveredAt, &d.CreatedAt, &d.UpdatedAt)
|
||||
return d, err
|
||||
}
|
||||
|
||||
func (s *NotificationService) deliver(ctx context.Context, channel NotificationChannel, delivery NotificationDelivery) error {
|
||||
body, _ := json.Marshal(map[string]any{"event_id": delivery.EventID, "event_type": delivery.EventType, "occurred_at": delivery.CreatedAt.UTC().Format(time.RFC3339Nano), "payload": json.RawMessage(delivery.Payload)})
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, channel.WebhookURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("X-Gateway-Event-ID", delivery.EventID)
|
||||
if len(channel.EncryptedSigningSecret) > 0 {
|
||||
secret, decryptErr := s.cipher.Decrypt(channel.EncryptedSigningSecret, channel.SigningSecretKEKVersion)
|
||||
if decryptErr != nil {
|
||||
return decryptErr
|
||||
}
|
||||
mac := hmac.New(sha256.New, secret)
|
||||
_, _ = mac.Write(body)
|
||||
request.Header.Set("X-Gateway-Signature", "sha256="+hex.EncodeToString(mac.Sum(nil)))
|
||||
}
|
||||
response, requestErr := s.client.Do(request)
|
||||
status := 0
|
||||
if response != nil {
|
||||
status = response.StatusCode
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64<<10))
|
||||
response.Body.Close()
|
||||
}
|
||||
success := requestErr == nil && status >= 200 && status < 300
|
||||
message := ""
|
||||
if requestErr != nil {
|
||||
message = requestErr.Error()
|
||||
} else if !success {
|
||||
message = fmt.Sprintf("Webhook 返回 HTTP %d", status)
|
||||
}
|
||||
if len(message) > 1000 {
|
||||
message = message[:1000]
|
||||
}
|
||||
_, dbErr := s.assets.pool.Exec(context.WithoutCancel(ctx), `UPDATE gateway.notification_deliveries SET status=$2,attempts=attempts+1,response_status=nullif($3,0),last_error=$4,delivered_at=CASE WHEN $2='delivered' THEN clock_timestamp() ELSE delivered_at END,updated_at=clock_timestamp() WHERE id=$1`, delivery.ID, map[bool]string{true: "delivered", false: "failed"}[success], status, message)
|
||||
if dbErr != nil {
|
||||
return dbErr
|
||||
}
|
||||
if !success {
|
||||
return errors.New(message)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *NotificationService) RetryDelivery(ctx context.Context, id string) error {
|
||||
var d NotificationDelivery
|
||||
err := s.assets.pool.QueryRow(ctx, deliverySelect+` WHERE d.id=$1`, id).Scan(&d.ID, &d.ChannelID, &d.ChannelName, &d.EventID, &d.EventType, &d.Status, &d.LastError, &d.Payload, &d.Attempts, &d.ResponseStatus, &d.DeliveredAt, &d.CreatedAt, &d.UpdatedAt)
|
||||
if err != nil {
|
||||
return mapNotFound(err)
|
||||
}
|
||||
channel, err := s.GetChannel(ctx, d.ChannelID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.deliver(ctx, channel, d)
|
||||
}
|
||||
|
||||
type NotificationDispatcher struct {
|
||||
service *NotificationService
|
||||
redis *redis.Client
|
||||
stream, group, consumer string
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func NewNotificationDispatcher(service *NotificationService, client *redis.Client, stream, consumer string, logger *slog.Logger) *NotificationDispatcher {
|
||||
return &NotificationDispatcher{service: service, redis: client, stream: stream, group: "gateway-notifications-v1", consumer: consumer, logger: logger}
|
||||
}
|
||||
func (d *NotificationDispatcher) Run(ctx context.Context) error {
|
||||
if err := d.redis.XGroupCreateMkStream(ctx, d.stream, d.group, "$").Err(); err != nil && !strings.Contains(err.Error(), "BUSYGROUP") {
|
||||
return err
|
||||
}
|
||||
delay := time.Duration(0)
|
||||
for {
|
||||
// Compensate for events claimed by a previous iteration (or a crashed
|
||||
// worker) that were never acknowledged. XReadGroup ">" never re-reads
|
||||
// the pending list, so without this pass those events would be silently
|
||||
// dropped for good.
|
||||
if err := d.reclaim(ctx); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
if d.logger != nil {
|
||||
d.logger.Error("notification pending reclaim failed; will retry", "error", err)
|
||||
}
|
||||
}
|
||||
streams, err := d.redis.XReadGroup(ctx, &redis.XReadGroupArgs{Group: d.group, Consumer: d.consumer, Streams: []string{d.stream, ">"}, Count: 20, Block: 5 * time.Second}).Result()
|
||||
if errors.Is(err, redis.Nil) {
|
||||
delay = 0
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
if d.logger != nil {
|
||||
d.logger.Error("notification stream read failed; backing off", "error", err)
|
||||
}
|
||||
// Reconnect with backoff instead of killing the worker: a transient
|
||||
// Redis blip must not take the notification worker down and let the
|
||||
// stream trim every queued event.
|
||||
if !d.wait(ctx, &delay) {
|
||||
return nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
delay = 0
|
||||
for _, stream := range streams {
|
||||
for _, message := range stream.Messages {
|
||||
if err = d.handle(ctx, message); err != nil {
|
||||
// Do not acknowledge: the event stays in the pending list and
|
||||
// is retried by the reclaim pass above.
|
||||
if d.logger != nil {
|
||||
d.logger.Error("notification event failed; will retry", "stream_id", message.ID, "error", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err = d.redis.XAck(ctx, d.stream, d.group, message.ID).Err(); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
// An ack failure must not drop a successfully handled event;
|
||||
// it is reclaimed and re-acked on the next pass (handle is
|
||||
// idempotent via the delivery table).
|
||||
if d.logger != nil {
|
||||
d.logger.Error("notification ack failed; will retry via reclaim", "stream_id", message.ID, "error", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// reclaim re-processes pending stream entries that have been idle longer than
|
||||
// the threshold. handle is idempotent (deliveries are keyed by channel+event),
|
||||
// so re-running it only creates the delivery rows that were never created or
|
||||
// updates attempts on ones already recorded.
|
||||
func (d *NotificationDispatcher) reclaim(ctx context.Context) error {
|
||||
for {
|
||||
messages, cursor, err := d.redis.XAutoClaim(ctx, &redis.XAutoClaimArgs{
|
||||
Stream: d.stream,
|
||||
Group: d.group,
|
||||
Consumer: d.consumer,
|
||||
MinIdle: 30 * time.Second,
|
||||
Start: "0",
|
||||
Count: 20,
|
||||
}).Result()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, message := range messages {
|
||||
if err = d.handle(ctx, message); err != nil {
|
||||
if d.logger != nil {
|
||||
d.logger.Warn("notification pending retry failed", "stream_id", message.ID, "error", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err = d.redis.XAck(ctx, d.stream, d.group, message.ID).Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if cursor == "0-0" {
|
||||
break
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// wait sleeps with exponential backoff, capping at 30s. It returns false when
|
||||
// the context is cancelled.
|
||||
func (d *NotificationDispatcher) wait(ctx context.Context, delay *time.Duration) bool {
|
||||
const max = 30 * time.Second
|
||||
if *delay == 0 {
|
||||
*delay = 200 * time.Millisecond
|
||||
} else {
|
||||
*delay *= 2
|
||||
if *delay > max {
|
||||
*delay = max
|
||||
}
|
||||
}
|
||||
timer := time.NewTimer(*delay)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case <-timer.C:
|
||||
return true
|
||||
}
|
||||
}
|
||||
func (d *NotificationDispatcher) handle(ctx context.Context, message redis.XMessage) error {
|
||||
eventID := fmt.Sprint(message.Values["event_id"])
|
||||
eventType := fmt.Sprint(message.Values["event_type"])
|
||||
payload := json.RawMessage(fmt.Sprint(message.Values["payload"]))
|
||||
channels, err := d.service.ListChannels(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, channel := range channels {
|
||||
if !channel.Enabled || !matchesEvent(channel.EventPatterns, eventType) {
|
||||
continue
|
||||
}
|
||||
delivery, deliveryErr := d.service.ensureDelivery(ctx, channel, eventID, eventType, payload)
|
||||
if deliveryErr != nil {
|
||||
return deliveryErr
|
||||
}
|
||||
if delivery.Status == "delivered" {
|
||||
continue
|
||||
}
|
||||
if deliveryErr = d.service.deliver(ctx, channel, delivery); deliveryErr != nil && d.logger != nil {
|
||||
d.logger.Warn("webhook delivery failed", "channel", channel.Name, "event_id", eventID, "error", deliveryErr)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type PromptCategory struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
TemplateCount int `json:"template_count"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
func (s *Service) ListPromptCategories(ctx context.Context) ([]PromptCategory, error) {
|
||||
rows, err := s.pool.Query(ctx, `SELECT c.id::text,c.name,c.description,count(t.id)::int,c.created_at FROM gateway.prompt_categories c LEFT JOIN gateway.prompt_templates t ON t.category_id=c.id GROUP BY c.id ORDER BY c.name`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []PromptCategory{}
|
||||
for rows.Next() {
|
||||
var item PromptCategory
|
||||
if err := rows.Scan(&item.ID, &item.Name, &item.Description, &item.TemplateCount, &item.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) CreatePromptCategory(ctx context.Context, name, description, actorID string) (PromptCategory, error) {
|
||||
name, description = strings.TrimSpace(name), strings.TrimSpace(description)
|
||||
if name == "" || len(name) > 64 || len(description) > 512 {
|
||||
return PromptCategory{}, errors.New("分类名称或描述格式无效")
|
||||
}
|
||||
id, err := newUUID()
|
||||
if err != nil {
|
||||
return PromptCategory{}, err
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return PromptCategory{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
var item PromptCategory
|
||||
err = tx.QueryRow(ctx, `INSERT INTO gateway.prompt_categories(id,name,description) VALUES($1,$2,$3) RETURNING id::text,name,description,created_at`, id, name, description).Scan(&item.ID, &item.Name, &item.Description, &item.CreatedAt)
|
||||
if err != nil {
|
||||
return PromptCategory{}, err
|
||||
}
|
||||
if err = emit(ctx, tx, "prompt_category.created", "prompt_category", id, actorID, nil); err != nil {
|
||||
return PromptCategory{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return PromptCategory{}, err
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (s *Service) DeletePromptCategory(ctx context.Context, id, actorID string) error {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
if _, err = tx.Exec(ctx, `UPDATE gateway.prompt_templates SET category_id=NULL,revision=revision+1,updated_at=clock_timestamp() WHERE category_id=$1`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.prompt_categories WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "prompt_category.deleted", "prompt_category", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func validateVariables(values []Variable) ([]Variable, error) {
|
||||
seen := map[string]struct{}{}
|
||||
for i := range values {
|
||||
values[i].Name = strings.TrimSpace(values[i].Name)
|
||||
values[i].Label = strings.TrimSpace(values[i].Label)
|
||||
if !regexpVariableName(values[i].Name) {
|
||||
return nil, fmt.Errorf("变量名 %q 格式无效", values[i].Name)
|
||||
}
|
||||
if _, ok := seen[values[i].Name]; ok {
|
||||
return nil, fmt.Errorf("变量 %q 重复", values[i].Name)
|
||||
}
|
||||
seen[values[i].Name] = struct{}{}
|
||||
if len(values[i].Default) > 10000 || len(values[i].Label) > 128 {
|
||||
return nil, errors.New("变量定义过长")
|
||||
}
|
||||
}
|
||||
if len(values) > 100 {
|
||||
return nil, errors.New("变量最多 100 个")
|
||||
}
|
||||
return values, nil
|
||||
}
|
||||
|
||||
func regexpVariableName(value string) bool {
|
||||
if value == "" {
|
||||
return false
|
||||
}
|
||||
matched := variableRE.FindStringSubmatch("{{" + value + "}}")
|
||||
return len(matched) == 2 && matched[1] == value
|
||||
}
|
||||
|
||||
func validatePromptInput(input *PromptInput, create bool) error {
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Description = strings.TrimSpace(input.Description)
|
||||
input.Content = strings.TrimSpace(input.Content)
|
||||
input.ChangeNote = strings.TrimSpace(input.ChangeNote)
|
||||
if input.Name == "" || len(input.Name) > 128 || len(input.Description) > 4000 {
|
||||
return errors.New("Prompt 名称或描述格式无效")
|
||||
}
|
||||
if create && (input.Content == "" || len(input.Content) > 100000) {
|
||||
return errors.New("Prompt 正文不能为空且最多 100000 字符")
|
||||
}
|
||||
var err error
|
||||
if input.Tags, err = normalizeStrings(input.Tags, 30); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.DepartmentIDs, err = normalizeStrings(input.DepartmentIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.Variables, err = validateVariables(input.Variables); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const promptSelect = `SELECT t.id::text,t.name,t.description,t.category_id::text,coalesce(c.name,''),t.tags,t.department_ids::text[],t.enabled,t.current_version,t.revision,t.created_at,t.updated_at,
|
||||
v.id::text,v.template_id::text,v.version,v.content,v.variables,v.change_note,v.created_at
|
||||
FROM gateway.prompt_templates t LEFT JOIN gateway.prompt_categories c ON c.id=t.category_id LEFT JOIN gateway.prompt_versions v ON v.template_id=t.id AND v.version=t.current_version`
|
||||
|
||||
func scanPrompt(row pgx.Row) (PromptTemplate, error) {
|
||||
var p PromptTemplate
|
||||
var vID, vTemplate *string
|
||||
var vVersion *int
|
||||
var content, change *string
|
||||
var variables []byte
|
||||
var versionCreated *time.Time
|
||||
err := row.Scan(&p.ID, &p.Name, &p.Description, &p.CategoryID, &p.CategoryName, &p.Tags, &p.DepartmentIDs, &p.Enabled, &p.CurrentVersion, &p.Revision, &p.CreatedAt, &p.UpdatedAt, &vID, &vTemplate, &vVersion, &content, &variables, &change, &versionCreated)
|
||||
if err != nil {
|
||||
return p, mapNotFound(err)
|
||||
}
|
||||
if vID != nil {
|
||||
v := PromptVersion{ID: *vID, TemplateID: *vTemplate, Version: *vVersion, Content: *content, ChangeNote: *change, CreatedAt: *versionCreated}
|
||||
_ = json.Unmarshal(variables, &v.Variables)
|
||||
p.Current = &v
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func (s *Service) ListPrompts(ctx context.Context) ([]PromptTemplate, error) {
|
||||
rows, err := s.pool.Query(ctx, promptSelect+` ORDER BY t.updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []PromptTemplate{}
|
||||
for rows.Next() {
|
||||
p, err := scanPrompt(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, p)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) GetPrompt(ctx context.Context, id string) (PromptTemplate, error) {
|
||||
return scanPrompt(s.pool.QueryRow(ctx, promptSelect+` WHERE t.id=$1`, id))
|
||||
}
|
||||
|
||||
func (s *Service) CreatePrompt(ctx context.Context, input PromptInput, actorID string) (PromptTemplate, error) {
|
||||
if err := validatePromptInput(&input, true); err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
if input.Variables == nil {
|
||||
input.Variables = []Variable{}
|
||||
}
|
||||
id, err := newUUID()
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
versionID, err := newUUID()
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
vars, _ := json.Marshal(input.Variables)
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.prompt_templates(id,name,description,category_id,tags,department_ids,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8)`, id, input.Name, input.Description, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled, actorID)
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.prompt_versions(id,template_id,version,content,variables,change_note,created_by) VALUES($1,$2,1,$3,$4,$5,$6)`, versionID, id, input.Content, vars, input.ChangeNote, actorID)
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `UPDATE gateway.prompt_templates SET current_version=1 WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
if err = emit(ctx, tx, "prompt.created", "prompt", id, actorID, map[string]any{"version": 1}); err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
return s.GetPrompt(ctx, id)
|
||||
}
|
||||
|
||||
func (s *Service) UpdatePrompt(ctx context.Context, id string, input PromptInput, actorID string) (PromptTemplate, error) {
|
||||
if err := validatePromptInput(&input, false); err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `UPDATE gateway.prompt_templates SET name=$2,description=$3,category_id=$4,tags=$5,department_ids=$6,enabled=$7,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Name, input.Description, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled)
|
||||
if err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return PromptTemplate{}, ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "prompt.updated", "prompt", id, actorID, nil); err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return PromptTemplate{}, err
|
||||
}
|
||||
return s.GetPrompt(ctx, id)
|
||||
}
|
||||
|
||||
func (s *Service) AddPromptVersion(ctx context.Context, id, content string, variables []Variable, changeNote, actorID string, activate bool) (PromptVersion, error) {
|
||||
content = strings.TrimSpace(content)
|
||||
changeNote = strings.TrimSpace(changeNote)
|
||||
if content == "" || len(content) > 100000 {
|
||||
return PromptVersion{}, errors.New("Prompt 正文不能为空且最多 100000 字符")
|
||||
}
|
||||
variables, err := validateVariables(variables)
|
||||
if err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
if variables == nil {
|
||||
variables = []Variable{}
|
||||
}
|
||||
raw, _ := json.Marshal(variables)
|
||||
versionID, err := newUUID()
|
||||
if err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
var lockedID string
|
||||
if err = tx.QueryRow(ctx, `SELECT id::text FROM gateway.prompt_templates WHERE id=$1 FOR UPDATE`, id).Scan(&lockedID); errors.Is(err, pgx.ErrNoRows) {
|
||||
return PromptVersion{}, ErrNotFound
|
||||
} else if err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
var version int
|
||||
err = tx.QueryRow(ctx, `SELECT coalesce(max(version),0)+1 FROM gateway.prompt_versions WHERE template_id=$1`, id).Scan(&version)
|
||||
if err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
var created time.Time
|
||||
err = tx.QueryRow(ctx, `INSERT INTO gateway.prompt_versions(id,template_id,version,content,variables,change_note,created_by) SELECT $1,id,$3,$4,$5,$6,$7 FROM gateway.prompt_templates WHERE id=$2 RETURNING created_at`, versionID, id, version, content, raw, changeNote, actorID).Scan(&created)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return PromptVersion{}, ErrNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
if activate {
|
||||
if _, err = tx.Exec(ctx, `UPDATE gateway.prompt_templates SET current_version=$2,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, version); err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
}
|
||||
if err = emit(ctx, tx, "prompt.version_created", "prompt", id, actorID, map[string]any{"version": version, "active": activate}); err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return PromptVersion{}, err
|
||||
}
|
||||
return PromptVersion{ID: versionID, TemplateID: id, Version: version, Content: content, Variables: variables, ChangeNote: changeNote, CreatedAt: created}, nil
|
||||
}
|
||||
|
||||
func (s *Service) ListPromptVersions(ctx context.Context, id string) ([]PromptVersion, error) {
|
||||
rows, err := s.pool.Query(ctx, `SELECT id::text,template_id::text,version,content,variables,change_note,created_at FROM gateway.prompt_versions WHERE template_id=$1 ORDER BY version DESC`, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []PromptVersion{}
|
||||
for rows.Next() {
|
||||
var v PromptVersion
|
||||
var raw []byte
|
||||
if err = rows.Scan(&v.ID, &v.TemplateID, &v.Version, &v.Content, &raw, &v.ChangeNote, &v.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = json.Unmarshal(raw, &v.Variables)
|
||||
items = append(items, v)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Service) ActivatePromptVersion(ctx context.Context, id string, version int, actorID string) error {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `UPDATE gateway.prompt_templates t SET current_version=$2,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1 AND EXISTS(SELECT 1 FROM gateway.prompt_versions v WHERE v.template_id=t.id AND v.version=$2)`, id, version)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "prompt.version_activated", "prompt", id, actorID, map[string]any{"version": version}); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *Service) DeletePrompt(ctx context.Context, id, actorID string) error {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.prompt_templates WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "prompt.deleted", "prompt", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
func (s *Service) RenderPromptByName(ctx context.Context, name string, provided map[string]any) (PromptVersion, string, error) {
|
||||
var v PromptVersion
|
||||
var raw []byte
|
||||
err := s.pool.QueryRow(ctx, `SELECT v.id::text,v.template_id::text,v.version,v.content,v.variables,v.change_note,v.created_at FROM gateway.prompt_templates t JOIN gateway.prompt_versions v ON v.template_id=t.id AND v.version=t.current_version WHERE t.name=$1 AND t.enabled`, name).Scan(&v.ID, &v.TemplateID, &v.Version, &v.Content, &raw, &v.ChangeNote, &v.CreatedAt)
|
||||
if err != nil {
|
||||
return v, "", mapNotFound(err)
|
||||
}
|
||||
_ = json.Unmarshal(raw, &v.Variables)
|
||||
rendered, err := RenderPrompt(v.Content, v.Variables, provided)
|
||||
return v, rendered, err
|
||||
}
|
||||
@@ -0,0 +1,688 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"aigateway.local/core/internal/factcheck"
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
type RuntimeHTTPHandler struct {
|
||||
service *Service
|
||||
tools *ToolService
|
||||
retriever Retriever
|
||||
auth apikey.PrincipalAuthenticator
|
||||
gateway http.Handler
|
||||
factCheck *factcheck.Engine
|
||||
logger *slog.Logger
|
||||
mux *http.ServeMux
|
||||
market MarketplaceDeps
|
||||
}
|
||||
|
||||
// MarketplaceDeps carries the resource-marketplace services into the runtime
|
||||
// handler (MCP servers, skills, digital employees, installations).
|
||||
type MarketplaceDeps struct {
|
||||
MCPServers *MCPServerService
|
||||
Skills *SkillService
|
||||
Employees *DigitalEmployeeService
|
||||
Market *MarketplaceService
|
||||
MCPClient *MCPClient
|
||||
}
|
||||
|
||||
func NewRuntimeHTTPHandler(service *Service, tools *ToolService, retriever Retriever, auth apikey.PrincipalAuthenticator, gatewayHandler http.Handler, market MarketplaceDeps) *RuntimeHTTPHandler {
|
||||
h := &RuntimeHTTPHandler{service: service, tools: tools, retriever: retriever, auth: auth, gateway: gatewayHandler, mux: http.NewServeMux(), market: market}
|
||||
h.logger = slog.Default()
|
||||
h.mux.HandleFunc("GET /v1/prompts", h.listPrompts)
|
||||
h.mux.HandleFunc("POST /v1/prompts/{name}/render", h.renderPrompt)
|
||||
h.mux.HandleFunc("POST /v1/knowledge/search", h.searchKnowledge)
|
||||
h.mux.HandleFunc("POST /v1/knowledge/{id}/search", h.searchKnowledge)
|
||||
h.mux.HandleFunc("GET /v1/tools", h.listTools)
|
||||
h.mux.HandleFunc("POST /v1/tools/{code}/invoke", h.invokeTool)
|
||||
h.mux.HandleFunc("POST /v1/applications/{code}/chat/completions", h.runApplication)
|
||||
h.mux.HandleFunc("POST /v1/skills/{code}/render", h.renderSkill)
|
||||
h.mux.HandleFunc("GET /v1/mcp-servers", h.listMCPServers)
|
||||
h.mux.HandleFunc("GET /v1/mcp-servers/{code}/tools", h.mcpServerTools)
|
||||
h.mux.HandleFunc("POST /v1/mcp-servers/{code}/tools/{tool}/invoke", h.invokeMCPTool)
|
||||
h.mux.HandleFunc("POST /v1/digital-employees/{code}/chat/completions", h.runDigitalEmployee)
|
||||
return h
|
||||
}
|
||||
func (h *RuntimeHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
|
||||
|
||||
// SetLogger wires a logger for best-effort diagnostics (fact-check skips etc.).
|
||||
func (h *RuntimeHTTPHandler) SetLogger(logger *slog.Logger) {
|
||||
if logger != nil {
|
||||
h.logger = logger
|
||||
}
|
||||
}
|
||||
|
||||
// SetFactCheckEngine enables post-answer fact-checking on application
|
||||
// conversations. When nil (the default) fact-checking is skipped entirely.
|
||||
func (h *RuntimeHTTPHandler) SetFactCheckEngine(engine *factcheck.Engine) { h.factCheck = engine }
|
||||
|
||||
// factCheckRetriever adapts the workbench Retriever to the fact-check engine's
|
||||
// EvidenceRetriever interface, reusing the same knowledge-base search path that
|
||||
// application prompts already use.
|
||||
type factCheckRetriever struct{ inner Retriever }
|
||||
|
||||
func NewFactCheckRetriever(inner Retriever) *factCheckRetriever {
|
||||
if inner == nil {
|
||||
return nil
|
||||
}
|
||||
return &factCheckRetriever{inner: inner}
|
||||
}
|
||||
|
||||
func (a *factCheckRetriever) Search(ctx context.Context, knowledgeBaseID, query string, topK int) ([]factcheck.EvidenceHit, error) {
|
||||
hits, err := a.inner.Search(ctx, knowledgeBaseID, query, topK)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]factcheck.EvidenceHit, 0, len(hits))
|
||||
for _, hit := range hits {
|
||||
out = append(out, factcheck.EvidenceHit{DocumentTitle: hit.DocumentTitle, Content: hit.Content})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// VerifyFactCheck satisfies factcheck.Verifier. It routes a non-streaming chat
|
||||
// completion through the same governed gateway, reusing the original request's
|
||||
// credential headers so the fact-check call is authenticated and rate-limited
|
||||
// exactly like the application call that produced the answer.
|
||||
func (h *RuntimeHTTPHandler) VerifyFactCheck(ctx context.Context, original *http.Request, model, system, user string, timeout time.Duration) (string, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
payload := map[string]any{
|
||||
"model": model,
|
||||
"temperature": 0,
|
||||
"stream": false,
|
||||
"messages": []map[string]any{
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user},
|
||||
},
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
request := original.Clone(ctx)
|
||||
request.Method = http.MethodPost
|
||||
request.URL.Path = "/v1/chat/completions"
|
||||
request.URL.RawPath = ""
|
||||
request.Body = ioNopCloser{bytes.NewReader(raw)}
|
||||
request.ContentLength = int64(len(raw))
|
||||
request.Header = request.Header.Clone()
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := newBoundedRecorder()
|
||||
h.gateway.ServeHTTP(recorder, request)
|
||||
if recorder.overrun > 0 {
|
||||
return "", errors.New("事实核查响应超过 2MB 上限")
|
||||
}
|
||||
result := recorder.Result()
|
||||
defer result.Body.Close()
|
||||
var decoded map[string]any
|
||||
if json.NewDecoder(result.Body).Decode(&decoded) != nil {
|
||||
return "", errors.New("事实核查响应无法解析")
|
||||
}
|
||||
if result.StatusCode < 200 || result.StatusCode >= 300 {
|
||||
return "", fmt.Errorf("事实核查模型调用失败(HTTP %d)", result.StatusCode)
|
||||
}
|
||||
content, _ := firstChoiceMessage(decoded)["content"].(string)
|
||||
if strings.TrimSpace(content) == "" {
|
||||
return "", errors.New("事实核查模型未返回文本")
|
||||
}
|
||||
return content, nil
|
||||
}
|
||||
func (h *RuntimeHTTPHandler) principal(w http.ResponseWriter, r *http.Request) (apikey.Principal, bool) {
|
||||
secret := strings.TrimSpace(r.Header.Get("X-Gateway-API-Key"))
|
||||
if secret == "" {
|
||||
value := strings.TrimSpace(r.Header.Get("Authorization"))
|
||||
if strings.HasPrefix(strings.ToLower(value), "bearer ") {
|
||||
secret = strings.TrimSpace(value[7:])
|
||||
}
|
||||
}
|
||||
principal, err := h.auth.AuthenticatePrincipal(r.Context(), secret)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, 401, "API Key 无效或已过期")
|
||||
return principal, false
|
||||
}
|
||||
return principal, true
|
||||
}
|
||||
func visible(departments []string, principal apikey.Principal, secure bool) bool {
|
||||
if principal.APIKeyID == "" {
|
||||
return true
|
||||
}
|
||||
if len(departments) == 0 {
|
||||
return !secure
|
||||
}
|
||||
if principal.TenantID == nil {
|
||||
return false
|
||||
}
|
||||
for _, id := range departments {
|
||||
if id == *principal.TenantID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) listPrompts(w http.ResponseWriter, r *http.Request) {
|
||||
principal, ok := h.principal(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.ListPrompts(r.Context())
|
||||
if err != nil {
|
||||
apiresponse.Error(w, 503, "Prompt 服务暂不可用")
|
||||
return
|
||||
}
|
||||
result := []map[string]any{}
|
||||
for _, item := range items {
|
||||
if item.Enabled && item.Current != nil && visible(item.DepartmentIDs, principal, false) {
|
||||
result = append(result, map[string]any{"name": item.Name, "description": item.Description, "tags": item.Tags, "version": item.Current.Version, "variables": item.Current.Variables})
|
||||
}
|
||||
}
|
||||
writeRuntime(w, 200, map[string]any{"object": "list", "data": result})
|
||||
}
|
||||
func (h *RuntimeHTTPHandler) renderPrompt(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
|
||||
}
|
||||
items, err := h.service.ListPrompts(r.Context())
|
||||
if err != nil {
|
||||
runtimeError(w, 503, "Prompt 服务暂不可用")
|
||||
return
|
||||
}
|
||||
var selected *PromptTemplate
|
||||
for i := range items {
|
||||
if items[i].Name == r.PathValue("name") && items[i].Enabled && visible(items[i].DepartmentIDs, principal, false) {
|
||||
selected = &items[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if selected == nil || selected.Current == nil {
|
||||
runtimeError(w, 404, "Prompt 不存在或不可见")
|
||||
return
|
||||
}
|
||||
rendered, err := RenderPrompt(selected.Current.Content, selected.Current.Variables, input.Variables)
|
||||
if err != nil {
|
||||
runtimeError(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
writeRuntime(w, 200, map[string]any{"name": selected.Name, "version": selected.Current.Version, "rendered": rendered})
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) searchKnowledge(w http.ResponseWriter, r *http.Request) {
|
||||
principal, ok := h.principal(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
KnowledgeBaseID string `json:"knowledge_base_id"`
|
||||
Query string `json:"query"`
|
||||
TopK int `json:"top_k"`
|
||||
}
|
||||
if !decodeRuntime(w, r, &input) {
|
||||
return
|
||||
}
|
||||
id := r.PathValue("id")
|
||||
if id == "" {
|
||||
id = input.KnowledgeBaseID
|
||||
}
|
||||
kb, err := h.service.GetKnowledgeBase(r.Context(), id)
|
||||
if err != nil || !kb.Enabled || !visible(kb.DepartmentIDs, principal, false) {
|
||||
runtimeError(w, 404, "知识库不存在或不可见")
|
||||
return
|
||||
}
|
||||
hits, err := h.retriever.Search(r.Context(), id, input.Query, input.TopK)
|
||||
if err != nil {
|
||||
runtimeError(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
writeRuntime(w, 200, map[string]any{"knowledge_base_id": id, "results": hits})
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) listTools(w http.ResponseWriter, r *http.Request) {
|
||||
principal, ok := h.principal(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.tools.List(r.Context())
|
||||
if err != nil {
|
||||
runtimeError(w, 503, "工具服务暂不可用")
|
||||
return
|
||||
}
|
||||
result := []map[string]any{}
|
||||
for _, tool := range items {
|
||||
if tool.Enabled && visible(tool.DepartmentIDs, principal, true) {
|
||||
result = append(result, map[string]any{"code": tool.Code, "name": tool.Name, "description": tool.Description, "input_schema": tool.InputSchema})
|
||||
}
|
||||
}
|
||||
writeRuntime(w, 200, map[string]any{"object": "list", "data": result})
|
||||
}
|
||||
func (h *RuntimeHTTPHandler) invokeTool(w http.ResponseWriter, r *http.Request) {
|
||||
principal, ok := h.principal(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
Input map[string]any `json:"input"`
|
||||
}
|
||||
if !decodeRuntime(w, r, &input) {
|
||||
return
|
||||
}
|
||||
tool, err := h.tools.GetByCode(r.Context(), r.PathValue("code"))
|
||||
if err != nil || !visible(tool.DepartmentIDs, principal, true) {
|
||||
runtimeError(w, 404, "工具不存在或不可调用")
|
||||
return
|
||||
}
|
||||
result, err := h.tools.Execute(r.Context(), tool, input.Input, principal.APIKeyID, gateway.RequestID(r.Context()))
|
||||
if err != nil {
|
||||
runtimeError(w, 502, err.Error())
|
||||
return
|
||||
}
|
||||
writeRuntime(w, 200, result)
|
||||
}
|
||||
|
||||
type applicationRequest struct {
|
||||
Messages []map[string]any `json:"messages"`
|
||||
Variables map[string]any `json:"variables"`
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) runApplication(w http.ResponseWriter, r *http.Request) {
|
||||
principal, ok := h.principal(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input applicationRequest
|
||||
if !decodeRuntime(w, r, &input) {
|
||||
return
|
||||
}
|
||||
app, err := h.service.GetPublishedApplicationByCode(r.Context(), r.PathValue("code"))
|
||||
if err != nil || app.PublishedConfig == nil || !visible(app.DepartmentIDs, principal, false) {
|
||||
runtimeError(w, 404, "应用不存在、未发布或不可见")
|
||||
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.application_runs(id,application_id,version,api_key_id,request_id,status,latency_ms,retrieval_count,tool_count,error) VALUES($1,$2,nullif($3,0),nullif($4,'')::uuid,$5,$6,$7,$8,$9,$10)`, runID, app.ID, valueOrZero(app.PublishedVersion), principal.APIKeyID, gateway.RequestID(r.Context()), status, time.Since(started).Milliseconds(), retrievalCount, toolCount, runError)
|
||||
}
|
||||
}()
|
||||
payload, toolsByCode, prepareErr := h.prepareApplication(r.Context(), app, input, principal, &retrievalCount)
|
||||
if prepareErr != nil {
|
||||
runError = prepareErr.Error()
|
||||
runtimeError(w, 400, runError)
|
||||
return
|
||||
}
|
||||
config := *app.PublishedConfig
|
||||
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 >= config.MaxToolRounds {
|
||||
runError = "工具调用轮次已达上限"
|
||||
runtimeError(w, 502, runError)
|
||||
return
|
||||
}
|
||||
choice := firstChoiceMessage(response)
|
||||
payload["messages"] = append(payload["messages"].([]map[string]any), choice)
|
||||
for _, call := range calls {
|
||||
tool, exists := toolsByCode[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.tools.Execute(r.Context(), tool, args, principal.APIKeyID, gateway.RequestID(r.Context()))
|
||||
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++
|
||||
}
|
||||
}
|
||||
if h.factCheck != nil {
|
||||
h.applyFactCheck(r, input, response)
|
||||
}
|
||||
response["application"] = map[string]any{"code": app.Code, "name": app.Name, "version": valueOrZero(app.PublishedVersion), "retrieval_count": retrievalCount, "tool_calls": toolCount}
|
||||
status = "success"
|
||||
copyHeaders(w.Header(), responseHeaders)
|
||||
writeRuntime(w, statusCode, response)
|
||||
}
|
||||
|
||||
// applyFactCheck verifies the assistant answer against configured knowledge
|
||||
// bases and applies the policy action. It must never fail the chat: any error
|
||||
// is logged and the answer is returned unchanged.
|
||||
func (h *RuntimeHTTPHandler) applyFactCheck(r *http.Request, input applicationRequest, response map[string]any) {
|
||||
answer, _ := assistantAnswer(response)
|
||||
lastQuestion := lastUserMessage(input.Messages)
|
||||
if strings.TrimSpace(answer) == "" || strings.TrimSpace(lastQuestion) == "" {
|
||||
return
|
||||
}
|
||||
verifier := func(ctx context.Context, model, system, user string, timeout time.Duration) (string, error) {
|
||||
return h.VerifyFactCheck(ctx, r, model, system, user, timeout)
|
||||
}
|
||||
event, err := h.factCheck.Check(r.Context(), gateway.RequestID(r.Context()), lastQuestion, answer, factcheck.VerifierFunc(verifier))
|
||||
if err != nil {
|
||||
h.logger.Warn("fact-check skipped", "request_id", gateway.RequestID(r.Context()), "error", err)
|
||||
return
|
||||
}
|
||||
if event.ID == "" {
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case event.Action == "block" && event.Verdict == "unsupported":
|
||||
overrideAnswer(response, "无法回答:该回复与知识库事实不符,已被事实核查拦截。")
|
||||
response["fact_check"] = map[string]any{"event_id": event.ID, "verdict": event.Verdict, "support_score": event.SupportScore, "blocked": true}
|
||||
case event.Action == "annotate":
|
||||
response["fact_check"] = map[string]any{"event_id": event.ID, "verdict": event.Verdict, "support_score": event.SupportScore, "blocked": false}
|
||||
}
|
||||
}
|
||||
|
||||
// lastUserMessage returns the content of the last user message in the request.
|
||||
func lastUserMessage(messages []map[string]any) string {
|
||||
last := ""
|
||||
for _, message := range messages {
|
||||
role, _ := message["role"].(string)
|
||||
if role != "user" {
|
||||
continue
|
||||
}
|
||||
if content, ok := message["content"].(string); ok {
|
||||
last = content
|
||||
}
|
||||
}
|
||||
return last
|
||||
}
|
||||
|
||||
// assistantAnswer extracts the final assistant text from a gateway response.
|
||||
func assistantAnswer(response map[string]any) (string, bool) {
|
||||
content, _ := firstChoiceMessage(response)["content"].(string)
|
||||
return content, strings.TrimSpace(content) != ""
|
||||
}
|
||||
|
||||
// overrideAnswer rewrites the assistant message content in place so the portal
|
||||
// and runtime consumers of response["choices"][0]["message"]["content"] all see
|
||||
// the fact-checked text.
|
||||
func overrideAnswer(response map[string]any, content string) {
|
||||
if message := firstChoiceMessage(response); message != nil {
|
||||
message["content"] = content
|
||||
}
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) prepareApplication(ctx context.Context, app Application, input applicationRequest, principal apikey.Principal, retrievalCount *int) (map[string]any, map[string]Tool, error) {
|
||||
config := *app.PublishedConfig
|
||||
messages := make([]map[string]any, 0, len(input.Messages)+2)
|
||||
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 config.PromptTemplateID != "" {
|
||||
prompt, err := h.service.GetPrompt(ctx, config.PromptTemplateID)
|
||||
// The prompt must be visible to this principal, mirroring the direct
|
||||
// render/list entry points, or a shared application could leak a
|
||||
// department-scoped prompt across departments.
|
||||
if err != nil || prompt.Current == nil || !prompt.Enabled || !visible(prompt.DepartmentIDs, principal, false) {
|
||||
return nil, nil, errors.New("应用绑定的 Prompt 当前不可用")
|
||||
}
|
||||
rendered, err := RenderPrompt(prompt.Current.Content, prompt.Current.Variables, input.Variables)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
system = append(system, rendered)
|
||||
}
|
||||
evidence := []string{}
|
||||
for _, kbID := range config.KnowledgeBaseIDs {
|
||||
kb, err := h.service.GetKnowledgeBase(ctx, kbID)
|
||||
if err != nil || !kb.Enabled || !visible(kb.DepartmentIDs, principal, false) {
|
||||
return nil, nil, fmt.Errorf("应用绑定的知识库 %s 当前不可用", kbID)
|
||||
}
|
||||
hits, searchErr := h.retriever.Search(ctx, kbID, lastQuestion, config.RetrievalTopK)
|
||||
if searchErr != nil {
|
||||
continue
|
||||
}
|
||||
for _, hit := range hits {
|
||||
*retrievalCount++
|
||||
evidence = append(evidence, fmt.Sprintf("[资料%d|%s]\n%s", len(evidence)+1, hit.DocumentTitle, hit.Content))
|
||||
}
|
||||
}
|
||||
if len(evidence) > 0 {
|
||||
system = append(system, "请优先依据以下企业资料回答;资料不足时明确说明不确定,不得编造。引用时使用[资料N]。\n\n"+strings.Join(evidence, "\n\n"))
|
||||
}
|
||||
if len(system) > 0 {
|
||||
messages = append([]map[string]any{{"role": "system", "content": strings.Join(system, "\n\n")}}, messages...)
|
||||
}
|
||||
toolsByCode := map[string]Tool{}
|
||||
schemas := []map[string]any{}
|
||||
for _, toolID := range config.ToolIDs {
|
||||
tool, err := h.tools.Get(ctx, toolID)
|
||||
// Enforce the same department visibility as the direct tool invoke
|
||||
// entry point (secure=true because tools carry embedded request
|
||||
// headers); otherwise an app shared across departments could trigger a
|
||||
// department-only tool and borrow its stored credentials.
|
||||
if err != nil || !tool.Enabled || !visible(tool.DepartmentIDs, principal, true) {
|
||||
return nil, nil, fmt.Errorf("应用绑定的工具 %s 当前不可用", toolID)
|
||||
}
|
||||
toolsByCode[tool.Code] = tool
|
||||
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}})
|
||||
}
|
||||
payload := map[string]any{"model": config.Model, "messages": messages, "stream": false, "temperature": config.Temperature}
|
||||
if len(schemas) > 0 {
|
||||
payload["tools"] = schemas
|
||||
payload["tool_choice"] = "auto"
|
||||
}
|
||||
return payload, toolsByCode, nil
|
||||
}
|
||||
|
||||
func (h *RuntimeHTTPHandler) callGateway(original *http.Request, payload map[string]any) (int, http.Header, map[string]any, error) {
|
||||
raw, _ := json.Marshal(payload)
|
||||
request := original.Clone(original.Context())
|
||||
request.Method = http.MethodPost
|
||||
request.URL.Path = "/v1/chat/completions"
|
||||
request.URL.RawPath = ""
|
||||
request.Body = http.NoBody
|
||||
if len(raw) > 0 {
|
||||
request.Body = ioNopCloser{bytes.NewReader(raw)}
|
||||
}
|
||||
request.ContentLength = int64(len(raw))
|
||||
request.Header = request.Header.Clone()
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := newBoundedRecorder()
|
||||
h.gateway.ServeHTTP(recorder, request)
|
||||
if recorder.overrun > 0 {
|
||||
return http.StatusBadGateway, recorder.Header(), nil,
|
||||
errors.New("模型响应超过 2MB 上限,已截断")
|
||||
}
|
||||
result := recorder.Result()
|
||||
defer result.Body.Close()
|
||||
var decoded map[string]any
|
||||
if json.NewDecoder(result.Body).Decode(&decoded) != nil {
|
||||
return result.StatusCode, result.Header, nil, errors.New("模型返回无法解析")
|
||||
}
|
||||
if result.StatusCode < 200 || result.StatusCode >= 300 {
|
||||
message := "应用模型调用失败"
|
||||
if value, ok := decoded["error"].(map[string]any); ok {
|
||||
if text, ok := value["message"].(string); ok {
|
||||
message = text
|
||||
}
|
||||
}
|
||||
return result.StatusCode, result.Header, decoded, errors.New(message)
|
||||
}
|
||||
return result.StatusCode, result.Header, decoded, nil
|
||||
}
|
||||
|
||||
type ioNopCloser struct{ *bytes.Reader }
|
||||
|
||||
func (ioNopCloser) Close() error { return nil }
|
||||
|
||||
// maxGatewayResponseBytes caps how much of a model response a buffered app
|
||||
// call may hold in memory. Non-streaming conversations go through a recorder
|
||||
// that buffers the full upstream response; without a cap a long completion
|
||||
// could exhaust process memory under concurrent app conversations.
|
||||
const maxGatewayResponseBytes = 2 << 20 // 2 MiB
|
||||
|
||||
// boundedRecorder is a minimal http.ResponseWriter that buffers the response
|
||||
// up to maxGatewayResponseBytes. Anything beyond the cap is discarded (but
|
||||
// counted) so a runaway upstream completion can never exhaust memory; callGateway
|
||||
// turns an overrun into an explicit error instead of decoding truncated JSON.
|
||||
type boundedRecorder struct {
|
||||
code int
|
||||
header http.Header
|
||||
body bytes.Buffer
|
||||
overrun int64
|
||||
}
|
||||
|
||||
func newBoundedRecorder() *boundedRecorder {
|
||||
return &boundedRecorder{code: http.StatusOK, header: make(http.Header)}
|
||||
}
|
||||
|
||||
func (r *boundedRecorder) Header() http.Header { return r.header }
|
||||
|
||||
func (r *boundedRecorder) WriteHeader(code int) {
|
||||
if r.code != 0 {
|
||||
return
|
||||
}
|
||||
r.code = code
|
||||
}
|
||||
|
||||
func (r *boundedRecorder) Write(data []byte) (int, error) {
|
||||
if r.code == 0 {
|
||||
r.code = http.StatusOK
|
||||
}
|
||||
remaining := int64(maxGatewayResponseBytes) - int64(r.body.Len())
|
||||
if remaining > 0 {
|
||||
written := data
|
||||
if int64(len(written)) > remaining {
|
||||
written = written[:remaining]
|
||||
}
|
||||
_, _ = r.body.Write(written)
|
||||
}
|
||||
r.overrun += int64(len(data)) - min(remaining, int64(len(data)))
|
||||
return len(data), nil
|
||||
}
|
||||
|
||||
func (r *boundedRecorder) Flush() {}
|
||||
|
||||
func (r *boundedRecorder) Result() *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: r.code,
|
||||
Header: r.header,
|
||||
Body: io.NopCloser(bytes.NewReader(r.body.Bytes())),
|
||||
}
|
||||
}
|
||||
|
||||
type toolCall struct{ ID, Name, Arguments string }
|
||||
|
||||
func extractToolCalls(response map[string]any) []toolCall {
|
||||
message := firstChoiceMessage(response)
|
||||
rawCalls, ok := message["tool_calls"].([]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
calls := []toolCall{}
|
||||
for _, raw := range rawCalls {
|
||||
item, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fn, _ := item["function"].(map[string]any)
|
||||
calls = append(calls, toolCall{ID: toString(item["id"]), Name: toString(fn["name"]), Arguments: toString(fn["arguments"])})
|
||||
}
|
||||
return calls
|
||||
}
|
||||
func firstChoiceMessage(response map[string]any) map[string]any {
|
||||
choices, ok := response["choices"].([]any)
|
||||
if !ok || len(choices) == 0 {
|
||||
return map[string]any{}
|
||||
}
|
||||
choice, _ := choices[0].(map[string]any)
|
||||
message, _ := choice["message"].(map[string]any)
|
||||
return message
|
||||
}
|
||||
func valueOrZero(value *int) int {
|
||||
if value == nil {
|
||||
return 0
|
||||
}
|
||||
return *value
|
||||
}
|
||||
func decodeRuntime(w http.ResponseWriter, r *http.Request, target any) bool {
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 2<<20))
|
||||
decoder.UseNumber()
|
||||
if err := decoder.Decode(target); err != nil {
|
||||
runtimeError(w, 400, "请求格式无效")
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
func runtimeError(w http.ResponseWriter, status int, message string) {
|
||||
writeRuntime(w, status, map[string]any{"error": map[string]any{"message": message, "type": "application_error"}})
|
||||
}
|
||||
func writeRuntime(w http.ResponseWriter, status int, value any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(value)
|
||||
}
|
||||
func copyHeaders(target, source http.Header) {
|
||||
for key, values := range source {
|
||||
lower := strings.ToLower(key)
|
||||
if lower == "connection" || lower == "content-length" || lower == "transfer-encoding" {
|
||||
continue
|
||||
}
|
||||
target[key] = append([]string(nil), values...)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,438 @@
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type Service struct{ pool *pgxpool.Pool }
|
||||
|
||||
func NewService(pool *pgxpool.Pool) *Service { return &Service{pool: pool} }
|
||||
|
||||
func newUUID() (string, error) { return platformid.NewUUID() }
|
||||
|
||||
func emit(ctx context.Context, tx pgx.Tx, eventType, aggregateType, aggregateID, actorID string, values map[string]any) error {
|
||||
eventID, err := newUUID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if values == nil {
|
||||
values = make(map[string]any)
|
||||
}
|
||||
values[aggregateType+"_id"] = aggregateID
|
||||
values["actor_id"] = actorID
|
||||
payload, err := json.Marshal(values)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,$3,$4,$5)`, eventID, eventType, aggregateType, aggregateID, payload)
|
||||
return err
|
||||
}
|
||||
|
||||
func normalizeStrings(values []string, maximum int) ([]string, error) {
|
||||
seen := make(map[string]struct{}, len(values))
|
||||
result := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
continue
|
||||
}
|
||||
if len(value) > 255 {
|
||||
return nil, errors.New("列表项过长")
|
||||
}
|
||||
if _, ok := seen[value]; ok {
|
||||
continue
|
||||
}
|
||||
seen[value] = struct{}{}
|
||||
result = append(result, value)
|
||||
}
|
||||
if len(result) > maximum {
|
||||
return nil, fmt.Errorf("列表最多允许 %d 项", maximum)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func mapNotFound(err error) error {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return ErrNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func rollback(ctx context.Context, tx pgx.Tx) { _ = tx.Rollback(ctx) }
|
||||
@@ -0,0 +1,223 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
type Skill struct {
|
||||
ID string `json:"id"`
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Content string `json:"content"`
|
||||
Variables []Variable `json:"variables"`
|
||||
ToolIDs []string `json:"tool_ids"`
|
||||
MCPServerIDs []string `json:"mcp_server_ids"`
|
||||
KnowledgeBaseIDs []string `json:"knowledge_base_ids"`
|
||||
Status string `json:"status"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
CategoryName string `json:"category_name"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type SkillInput struct {
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Content string `json:"content"`
|
||||
Variables []Variable `json:"variables"`
|
||||
ToolIDs []string `json:"tool_ids"`
|
||||
MCPServerIDs []string `json:"mcp_server_ids"`
|
||||
KnowledgeBaseIDs []string `json:"knowledge_base_ids"`
|
||||
Status string `json:"status"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
// SkillService manages packaged capabilities (skills): prompt content plus
|
||||
// optional tool / MCP server / knowledge-base bindings, publishable to the
|
||||
// marketplace and consumed by digital employees.
|
||||
type SkillService struct {
|
||||
assets *Service
|
||||
}
|
||||
|
||||
func NewSkillService(assets *Service) *SkillService { return &SkillService{assets: assets} }
|
||||
|
||||
func (s *SkillService) validate(ctx context.Context, input *SkillInput, create bool) error {
|
||||
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
|
||||
input.Name = strings.TrimSpace(input.Name)
|
||||
input.Description = strings.TrimSpace(input.Description)
|
||||
input.Content = strings.TrimSpace(input.Content)
|
||||
if !codeRE.MatchString(input.Code) || input.Name == "" || len(input.Name) > 128 || len(input.Description) > 4000 {
|
||||
return errors.New("Skill 编码、名称或描述格式无效")
|
||||
}
|
||||
if create && (input.Content == "" || len(input.Content) > 100000) {
|
||||
return errors.New("Skill 正文不能为空且最多 100000 字符")
|
||||
}
|
||||
switch input.Status {
|
||||
case "", "draft":
|
||||
input.Status = "draft"
|
||||
case "published", "archived":
|
||||
default:
|
||||
return errors.New("无效的资源状态")
|
||||
}
|
||||
var err error
|
||||
if input.Variables, err = validateVariables(input.Variables); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.ToolIDs, err = normalizeStrings(input.ToolIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.MCPServerIDs, err = normalizeStrings(input.MCPServerIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.KnowledgeBaseIDs, err = normalizeStrings(input.KnowledgeBaseIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.CategoryID != nil && strings.TrimSpace(*input.CategoryID) == "" {
|
||||
input.CategoryID = nil
|
||||
}
|
||||
var categoryErr error
|
||||
input.CategoryID, categoryErr = validCategoryID(ctx, s.assets.pool, input.CategoryID, "skill")
|
||||
if categoryErr != nil {
|
||||
return categoryErr
|
||||
}
|
||||
if input.Tags, err = normalizeStrings(input.Tags, 30); err != nil {
|
||||
return err
|
||||
}
|
||||
if input.DepartmentIDs, err = normalizeStrings(input.DepartmentIDs, 100); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const skillSelect = `SELECT s.id::text,s.code,s.name,s.description,s.content,s.variables,s.tool_ids,s.mcp_server_ids,s.knowledge_base_ids,s.status,s.category_id::text,coalesce(c.name,''),s.tags,s.department_ids::text[],s.enabled,s.revision,s.created_at,s.updated_at FROM gateway.skills s LEFT JOIN gateway.marketplace_categories c ON c.id=s.category_id`
|
||||
|
||||
func scanSkill(row pgx.Row) (Skill, error) {
|
||||
var s Skill
|
||||
var variables []byte
|
||||
err := row.Scan(&s.ID, &s.Code, &s.Name, &s.Description, &s.Content, &variables, &s.ToolIDs, &s.MCPServerIDs, &s.KnowledgeBaseIDs, &s.Status, &s.CategoryID, &s.CategoryName, &s.Tags, &s.DepartmentIDs, &s.Enabled, &s.Revision, &s.CreatedAt, &s.UpdatedAt)
|
||||
if err != nil {
|
||||
return s, mapNotFound(err)
|
||||
}
|
||||
_ = json.Unmarshal(variables, &s.Variables)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *SkillService) List(ctx context.Context) ([]Skill, error) {
|
||||
rows, err := s.assets.pool.Query(ctx, skillSelect+` ORDER BY s.updated_at DESC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Skill{}
|
||||
for rows.Next() {
|
||||
skill, err := scanSkill(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, skill)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *SkillService) Get(ctx context.Context, id string) (Skill, error) {
|
||||
return scanSkill(s.assets.pool.QueryRow(ctx, skillSelect+` WHERE s.id=$1`, id))
|
||||
}
|
||||
|
||||
func (s *SkillService) GetByCode(ctx context.Context, code string) (Skill, error) {
|
||||
return scanSkill(s.assets.pool.QueryRow(ctx, skillSelect+` WHERE s.code=$1`, code))
|
||||
}
|
||||
|
||||
// GetPublishedByCode returns a published, enabled skill by code.
|
||||
func (s *SkillService) GetPublishedByCode(ctx context.Context, code string) (Skill, error) {
|
||||
return scanSkill(s.assets.pool.QueryRow(ctx, skillSelect+` WHERE s.code=$1 AND s.status='published' AND s.enabled`, code))
|
||||
}
|
||||
|
||||
func (s *SkillService) Save(ctx context.Context, id string, input SkillInput, actorID string, create bool) (Skill, error) {
|
||||
if err := s.validate(ctx, &input, create); err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
tx, err := s.assets.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
defer rollback(ctx, tx)
|
||||
if create {
|
||||
id, err = newUUID()
|
||||
if err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
variables, _ := json.Marshal(input.Variables)
|
||||
_, err = tx.Exec(ctx, `INSERT INTO gateway.skills(id,code,name,description,content,variables,tool_ids,mcp_server_ids,knowledge_base_ids,status,category_id,tags,department_ids,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.Content, variables, input.ToolIDs, input.MCPServerIDs, input.KnowledgeBaseIDs, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled, actorID)
|
||||
} else {
|
||||
variables, _ := json.Marshal(input.Variables)
|
||||
tag, updateErr := tx.Exec(ctx, `UPDATE gateway.skills SET code=$2,name=$3,description=$4,content=$5,variables=$6,tool_ids=$7,mcp_server_ids=$8,knowledge_base_ids=$9,status=$10,category_id=$11,tags=$12,department_ids=$13,enabled=$14,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.Content, variables, input.ToolIDs, input.MCPServerIDs, input.KnowledgeBaseIDs, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled)
|
||||
err = updateErr
|
||||
if err == nil && tag.RowsAffected() == 0 {
|
||||
return Skill{}, ErrNotFound
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
event := "skill.updated"
|
||||
if create {
|
||||
event = "skill.created"
|
||||
}
|
||||
if err = emit(ctx, tx, event, "skill", id, actorID, nil); err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
if err = tx.Commit(ctx); err != nil {
|
||||
return Skill{}, err
|
||||
}
|
||||
return s.Get(ctx, id)
|
||||
}
|
||||
|
||||
func (s *SkillService) 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.digital_employees WHERE $1 = ANY(skill_ids)
|
||||
UNION ALL SELECT 1 FROM gateway.marketplace_installations WHERE resource_type='skill' AND resource_id=$1
|
||||
)`, id).Scan(&used); err != nil {
|
||||
return err
|
||||
}
|
||||
if used {
|
||||
return ErrConflict
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `DELETE FROM gateway.skills WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
if err = emit(ctx, tx, "skill.deleted", "skill", id, actorID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
// Render applies the skill's variables to its content (same substitution as
|
||||
// prompt rendering).
|
||||
func (s *SkillService) Render(skill Skill, provided map[string]any) (string, error) {
|
||||
return RenderPrompt(skill.Content, skill.Variables, provided)
|
||||
}
|
||||
@@ -0,0 +1,365 @@
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrNotFound = errors.New("AI asset not found")
|
||||
ErrConflict = errors.New("AI asset conflict")
|
||||
codeRE = regexp.MustCompile(`^[a-z][a-z0-9_-]{1,63}$`)
|
||||
variableRE = regexp.MustCompile(`\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}`)
|
||||
)
|
||||
|
||||
type Variable struct {
|
||||
Name string `json:"name"`
|
||||
Label string `json:"label,omitempty"`
|
||||
Default string `json:"default,omitempty"`
|
||||
Required bool `json:"required"`
|
||||
}
|
||||
|
||||
type PromptTemplate struct {
|
||||
ID, Name, Description string `json:"-"`
|
||||
CategoryID *string `json:"-"`
|
||||
CategoryName string `json:"-"`
|
||||
Tags []string `json:"-"`
|
||||
DepartmentIDs []string `json:"-"`
|
||||
Enabled bool `json:"-"`
|
||||
CurrentVersion *int `json:"-"`
|
||||
Current *PromptVersion `json:"-"`
|
||||
Revision int64 `json:"-"`
|
||||
CreatedAt, UpdatedAt time.Time `json:"-"`
|
||||
}
|
||||
|
||||
func (p PromptTemplate) MarshalJSON() ([]byte, error) {
|
||||
type output struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
CategoryID *string `json:"category_id,omitempty"`
|
||||
CategoryName string `json:"category_name"`
|
||||
Tags []string `json:"tags"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
CurrentVersion *int `json:"current_version,omitempty"`
|
||||
Current *PromptVersion `json:"current,omitempty"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
return json.Marshal(output{p.ID, p.Name, p.Description, p.CategoryID, p.CategoryName, nonNil(p.Tags), nonNil(p.DepartmentIDs), p.Enabled, p.CurrentVersion, p.Current, p.Revision, p.CreatedAt, p.UpdatedAt})
|
||||
}
|
||||
|
||||
type PromptVersion struct {
|
||||
ID string `json:"id"`
|
||||
TemplateID string `json:"template_id"`
|
||||
Version int `json:"version"`
|
||||
Content string `json:"content"`
|
||||
Variables []Variable `json:"variables"`
|
||||
ChangeNote string `json:"change_note"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type PromptInput struct {
|
||||
Name, Description string
|
||||
CategoryID *string
|
||||
Tags, DepartmentIDs []string
|
||||
Enabled bool
|
||||
Content, ChangeNote string
|
||||
Variables []Variable
|
||||
}
|
||||
|
||||
type KnowledgeBase struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
RetrievalMode string `json:"retrieval_mode"`
|
||||
ChunkSize int `json:"chunk_size"`
|
||||
ChunkOverlap int `json:"chunk_overlap"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Revision int64 `json:"revision"`
|
||||
DocumentCount int `json:"document_count"`
|
||||
ChunkCount int `json:"chunk_count"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type KnowledgeDocument struct {
|
||||
ID string `json:"id"`
|
||||
KnowledgeBaseID string `json:"knowledge_base_id"`
|
||||
Title string `json:"title"`
|
||||
SourceType string `json:"source_type"`
|
||||
SourceURI string `json:"source_uri"`
|
||||
ContentSHA256 string `json:"content_sha256"`
|
||||
Status string `json:"status"`
|
||||
StatusMessage string `json:"status_message"`
|
||||
CharCount int `json:"char_count"`
|
||||
ChunkCount int `json:"chunk_count"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type SearchHit struct {
|
||||
ChunkID string `json:"chunk_id"`
|
||||
DocumentID string `json:"document_id"`
|
||||
DocumentTitle string `json:"document_title"`
|
||||
Content string `json:"content"`
|
||||
ChunkIndex int `json:"chunk_index"`
|
||||
Score float64 `json:"score"`
|
||||
}
|
||||
|
||||
type Tool struct {
|
||||
ID string `json:"id"`
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
EndpointURL string `json:"endpoint_url"`
|
||||
HTTPMethod string `json:"http_method"`
|
||||
InputSchema json.RawMessage `json:"input_schema"`
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
Enabled bool `json:"enabled"`
|
||||
HasSecretHeaders bool `json:"has_secret_headers"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
EncryptedHeaders []byte `json:"-"`
|
||||
HeadersKEKVersion int `json:"-"`
|
||||
}
|
||||
|
||||
type ToolInput struct {
|
||||
Code, Name, Description, EndpointURL, HTTPMethod string
|
||||
Headers map[string]string
|
||||
InputSchema json.RawMessage
|
||||
TimeoutSeconds int
|
||||
DepartmentIDs []string
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
type ApplicationConfig struct {
|
||||
Model string `json:"model"`
|
||||
PromptTemplateID string `json:"prompt_template_id,omitempty"`
|
||||
KnowledgeBaseIDs []string `json:"knowledge_base_ids"`
|
||||
ToolIDs []string `json:"tool_ids"`
|
||||
RetrievalTopK int `json:"retrieval_top_k"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
MaxToolRounds int `json:"max_tool_rounds"`
|
||||
}
|
||||
|
||||
type Application struct {
|
||||
ID string `json:"id"`
|
||||
Code string `json:"code"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Status string `json:"status"`
|
||||
DepartmentIDs []string `json:"department_ids"`
|
||||
DraftConfig ApplicationConfig `json:"draft_config"`
|
||||
PublishedVersion *int `json:"published_version,omitempty"`
|
||||
PublishedConfig *ApplicationConfig `json:"published_config,omitempty"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type ApplicationRun struct {
|
||||
ID string `json:"id"`
|
||||
ApplicationID string `json:"application_id"`
|
||||
RequestID string `json:"request_id"`
|
||||
Status string `json:"status"`
|
||||
Error string `json:"error"`
|
||||
Version int `json:"version"`
|
||||
LatencyMS int64 `json:"latency_ms"`
|
||||
RetrievalCount int `json:"retrieval_count"`
|
||||
ToolCount int `json:"tool_count"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type NotificationChannel struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
WebhookURL string `json:"webhook_url"`
|
||||
EventPatterns []string `json:"event_patterns"`
|
||||
Enabled bool `json:"enabled"`
|
||||
HasSigningSecret bool `json:"has_signing_secret"`
|
||||
Revision int64 `json:"revision"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
EncryptedSigningSecret []byte `json:"-"`
|
||||
SigningSecretKEKVersion int `json:"-"`
|
||||
}
|
||||
|
||||
type NotificationDelivery struct {
|
||||
ID string `json:"id"`
|
||||
ChannelID string `json:"channel_id"`
|
||||
ChannelName string `json:"channel_name"`
|
||||
EventID string `json:"event_id"`
|
||||
EventType string `json:"event_type"`
|
||||
Status string `json:"status"`
|
||||
LastError string `json:"last_error"`
|
||||
Payload json.RawMessage `json:"payload"`
|
||||
Attempts int `json:"attempts"`
|
||||
ResponseStatus *int `json:"response_status,omitempty"`
|
||||
DeliveredAt *time.Time `json:"delivered_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
func RenderPrompt(content string, definitions []Variable, provided map[string]any) (string, error) {
|
||||
values := make(map[string]string, len(definitions))
|
||||
missing := make([]string, 0)
|
||||
for _, variable := range definitions {
|
||||
if raw, ok := provided[variable.Name]; ok && strings.TrimSpace(toString(raw)) != "" {
|
||||
values[variable.Name] = toString(raw)
|
||||
} else if variable.Default != "" {
|
||||
values[variable.Name] = variable.Default
|
||||
} else if variable.Required {
|
||||
missing = append(missing, variable.Name)
|
||||
} else {
|
||||
values[variable.Name] = ""
|
||||
}
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
return "", errors.New("缺少必填变量: " + strings.Join(missing, ", "))
|
||||
}
|
||||
return variableRE.ReplaceAllStringFunc(content, func(match string) string {
|
||||
parts := variableRE.FindStringSubmatch(match)
|
||||
if value, ok := values[parts[1]]; ok {
|
||||
return value
|
||||
}
|
||||
return match
|
||||
}), nil
|
||||
}
|
||||
|
||||
func toString(value any) string {
|
||||
if text, ok := value.(string); ok {
|
||||
return text
|
||||
}
|
||||
encoded, _ := json.Marshal(value)
|
||||
return string(encoded)
|
||||
}
|
||||
|
||||
func nonNil[T any](value []T) []T {
|
||||
if value == nil {
|
||||
return []T{}
|
||||
}
|
||||
return value
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"aigateway.local/core/internal/apikey"
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/cryptox"
|
||||
"aigateway.local/core/internal/platform/database"
|
||||
)
|
||||
|
||||
func TestWorkbenchPostgreSQLLifecycle(t *testing.T) {
|
||||
databaseURL := os.Getenv("WORKBENCH_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("WORKBENCH_TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8, MinConns: 0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
actorID := "11111111-1111-4111-8111-111111111111"
|
||||
_, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role) VALUES($1,'m4-test','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cleanup := func() {
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM gateway.notification_channels WHERE name='m4-webhook'; DELETE FROM gateway.applications WHERE code='m4_app'; DELETE FROM gateway.tool_definitions WHERE code='m4_lookup'; DELETE FROM gateway.knowledge_bases WHERE name='m4-integration-kb'; DELETE FROM gateway.prompt_templates WHERE name='m4-integration-prompt'`)
|
||||
}
|
||||
cleanup()
|
||||
defer cleanup()
|
||||
assets := NewService(pool)
|
||||
prompt, err := assets.CreatePrompt(ctx, PromptInput{Name: "m4-integration-prompt", Description: "test", Enabled: true, Content: "请回答 {{question}}", Variables: []Variable{{Name: "question", Required: true}}}, actorID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if prompt.Current == nil || prompt.Current.Version != 1 {
|
||||
t.Fatalf("unexpected prompt: %#v", prompt)
|
||||
}
|
||||
version, err := assets.AddPromptVersion(ctx, prompt.ID, "新版 {{question}}", []Variable{{Name: "question", Required: true}}, "v2", actorID, true)
|
||||
if err != nil || version.Version != 2 {
|
||||
t.Fatalf("version=%#v err=%v", version, err)
|
||||
}
|
||||
_, rendered, err := assets.RenderPromptByName(ctx, prompt.Name, map[string]any{"question": "可扩展吗"})
|
||||
if err != nil || rendered != "新版 可扩展吗" {
|
||||
t.Fatalf("render=%q err=%v", rendered, err)
|
||||
}
|
||||
kb, err := assets.SaveKnowledgeBase(ctx, KnowledgeBase{Name: "m4-integration-kb", Description: "test", RetrievalMode: "postgres_fts", ChunkSize: 200, ChunkOverlap: 20, Enabled: true}, actorID, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
doc, err := assets.AddKnowledgeDocument(ctx, kb.ID, "Go 架构说明", "text", "", "Go 网关采用不可变运行时快照。\n\nPostgreSQL 是权威配置存储。", actorID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if doc.ChunkCount == 0 {
|
||||
t.Fatal("expected chunks")
|
||||
}
|
||||
hits, err := NewPostgreSQLRetriever(assets).Search(ctx, kb.ID, "不可变运行时快照", 4)
|
||||
if err != nil || len(hits) == 0 {
|
||||
t.Fatalf("hits=%#v err=%v", hits, err)
|
||||
}
|
||||
key := base64.StdEncoding.EncodeToString(make([]byte, 32))
|
||||
cipher, err := cryptox.NewKeyring(key, 1, "", "m4-test-tools")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
receivedAuth := ""
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedAuth = r.Header.Get("Authorization")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"ok": true})
|
||||
}))
|
||||
defer upstream.Close()
|
||||
tools := NewToolService(assets, cipher, true)
|
||||
tool, err := tools.Save(ctx, "", ToolInput{Code: "m4_lookup", Name: "查询", EndpointURL: upstream.URL, HTTPMethod: "POST", Headers: map[string]string{"Authorization": "Bearer secret"}, InputSchema: json.RawMessage(`{"type":"object","required":["q"],"properties":{"q":{"type":"string"}}}`), TimeoutSeconds: 5, Enabled: true}, actorID, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result, err := tools.Execute(ctx, tool, map[string]any{"q": "gateway"}, "", "m4-test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result["status_code"] != http.StatusOK || receivedAuth != "Bearer secret" {
|
||||
t.Fatalf("result=%#v auth=%q", result, receivedAuth)
|
||||
}
|
||||
app, err := assets.SaveApplication(ctx, Application{Code: "m4_app", Name: "M4 App", Status: "draft", DraftConfig: ApplicationConfig{Model: "test-model", PromptTemplateID: prompt.ID, KnowledgeBaseIDs: []string{kb.ID}, ToolIDs: []string{tool.ID}, RetrievalTopK: 4, Temperature: .2, MaxToolRounds: 2}}, actorID, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
published, err := assets.PublishApplication(ctx, app.ID, "first", actorID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if published.PublishedVersion == nil || *published.PublishedVersion != 1 || published.Status != "active" {
|
||||
t.Fatalf("unexpected published app: %#v", published)
|
||||
}
|
||||
governed := false
|
||||
fakeGateway := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var payload map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
messages, _ := payload["messages"].([]any)
|
||||
if len(messages) > 0 && strings.Contains(toString(messages[0]), "不可变运行时快照") {
|
||||
governed = true
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "完成"}}}})
|
||||
})
|
||||
runtime := NewRuntimeHTTPHandler(assets, tools, NewPostgreSQLRetriever(assets), staticPrincipalAuthenticator{}, fakeGateway, MarketplaceDeps{})
|
||||
runtimeRequest := httptest.NewRequest(http.MethodPost, "/v1/applications/m4_app/chat/completions", bytes.NewBufferString(`{"messages":[{"role":"user","content":"不可变运行时快照是什么?"}],"variables":{"question":"架构"}}`))
|
||||
runtimeRequest.Header.Set("Authorization", "Bearer test")
|
||||
runtimeRequest = runtimeRequest.WithContext(gateway.WithRequestID(runtimeRequest.Context(), "m4-runtime"))
|
||||
runtimeResponse := httptest.NewRecorder()
|
||||
runtime.ServeHTTP(runtimeResponse, runtimeRequest)
|
||||
if runtimeResponse.Code != http.StatusOK || !governed || !strings.Contains(runtimeResponse.Body.String(), `"application"`) {
|
||||
t.Fatalf("runtime status=%d governed=%v body=%s", runtimeResponse.Code, governed, runtimeResponse.Body.String())
|
||||
}
|
||||
|
||||
signed := ""
|
||||
webhook := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
signed = r.Header.Get("X-Gateway-Signature")
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}))
|
||||
defer webhook.Close()
|
||||
notifications := NewNotificationService(assets, cipher, true)
|
||||
secret := "m4-signing-secret"
|
||||
channel, err := notifications.SaveChannel(ctx, "", NotificationInput{Name: "m4-webhook", WebhookURL: webhook.URL, SigningSecret: &secret, EventPatterns: []string{"application.*"}, Enabled: true}, actorID, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
eventID, _ := newUUID()
|
||||
delivery, err := notifications.ensureDelivery(ctx, channel, eventID, "application.published", json.RawMessage(`{"version":1}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = notifications.deliver(ctx, channel, delivery); err != nil || !strings.HasPrefix(signed, "sha256=") {
|
||||
t.Fatalf("delivery err=%v signature=%q", err, signed)
|
||||
}
|
||||
}
|
||||
|
||||
type staticPrincipalAuthenticator struct{}
|
||||
|
||||
func (staticPrincipalAuthenticator) AuthenticatePrincipal(context.Context, string) (apikey.Principal, error) {
|
||||
return apikey.Principal{}, nil
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package workbench
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRenderPromptValidatesRequiredVariables(t *testing.T) {
|
||||
definitions := []Variable{{Name: "topic", Required: true}, {Name: "tone", Default: "简洁"}}
|
||||
rendered, err := RenderPrompt("用{{ tone }}风格介绍 {{topic}},保留 {{unknown}}。", definitions, map[string]any{"topic": "Go"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rendered != "用简洁风格介绍 Go,保留 {{unknown}}。" {
|
||||
t.Fatalf("unexpected render: %q", rendered)
|
||||
}
|
||||
if _, err = RenderPrompt("{{topic}}", definitions, nil); err == nil {
|
||||
t.Fatal("expected missing variable error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkTextParagraphAwareAndOverlapped(t *testing.T) {
|
||||
text := "第一段。\n\n" + repeatRune('甲', 240) + "\n\n最后一段。"
|
||||
chunks := ChunkText(text, 200, 20)
|
||||
if len(chunks) < 3 {
|
||||
t.Fatalf("expected at least 3 chunks, got %d", len(chunks))
|
||||
}
|
||||
for _, chunk := range chunks {
|
||||
if RuneCount(chunk) > 222 {
|
||||
t.Fatalf("chunk unexpectedly large: %d", RuneCount(chunk))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeStringsStableDeduplication(t *testing.T) {
|
||||
values, err := normalizeStrings([]string{" b ", "a", "b", ""}, 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(values, []string{"b", "a"}) {
|
||||
t.Fatalf("unexpected values: %#v", values)
|
||||
}
|
||||
}
|
||||
|
||||
func repeatRune(value rune, count int) string {
|
||||
runes := make([]rune, count)
|
||||
for i := range runes {
|
||||
runes[i] = value
|
||||
}
|
||||
return string(runes)
|
||||
}
|
||||
Reference in New Issue
Block a user