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:
ben
2026-08-12 11:45:54 +08:00
commit 5759c1862e
807 changed files with 114727 additions and 0 deletions
+783
View File
@@ -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
+435
View File
@@ -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)
}
+380
View File
@@ -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, &note, &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
}
+378
View File
@@ -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()
}
+448
View File
@@ -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)
+444
View File
@@ -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)
}
}
+253
View File
@@ -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
}
+369
View File
@@ -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
}
+237
View File
@@ -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)
}
}
+437
View File
@@ -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
}
+377
View File
@@ -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
}
+688
View File
@@ -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...)
}
}
+438
View File
@@ -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)
}
+69
View File
@@ -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) }
+223
View File
@@ -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)
}
+365
View File
@@ -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))
}
}
+251
View File
@@ -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
}
+51
View File
@@ -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)
}