rogee/gochat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(captain): gate custom tools

Browse Source
This commit is contained in:
rogee
2026-06-07 15:50:23 +08:00
parent 0aff792b40
commit 0d8de3be0b
7 changed files with 81 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/handler/api/v1/captain_custom_tool_crud_handler_test.go
+21 -2
View File
@@ -51,12 +51,13 @@ func (s *CaptainCustomToolCRUDTestSuite) SetupSuite() {
)
s.Require().NoError(err)
account := &model.Account{Name: "CustomToolTestOrg", Locale: "en", Active: true}
account := &model.Account{Name: "CustomToolTestOrg", Locale: "en", Active: true, FeatureFlags: `{"custom_tools":true}`}
s.Require().NoError(db.Create(account).Error)
s.account = account
toolRepo := repository.NewCaptainCustomToolRepo(db)
svc := service.NewCaptainCustomToolService(toolRepo)
accountRepo := repository.NewAccountRepo(db)
svc := service.NewCaptainCustomToolService(toolRepo, accountRepo)
s.handler = NewCaptainCustomToolHandler(svc)
// Unified router: :id = account_id, :tool_id = tool_id
@@ -180,6 +181,24 @@ func (s *CaptainCustomToolCRUDTestSuite) TestCreate_无效accountID返回400() {
assert.Equal(s.T(), http.StatusBadRequest, w.Code)
}
func (s *CaptainCustomToolCRUDTestSuite) TestCreate_FeatureDisabledReturnsForbidden() {
disabledAccount := &model.Account{Name: "CustomToolDisabledOrg", Locale: "en", Active: true}
s.Require().NoError(s.db.Create(disabledAccount).Error)
path := "/api/v1/accounts/" + strconv.FormatUint(uint64(disabledAccount.ID), 10) + "/captain/custom_tools/"
w := s.makeRequest("POST", path, map[string]interface{}{
"custom_tool": map[string]interface{}{
"title": "blocked",
"endpoint_url": "https://example.com/blocked",
},
})
assert.Equal(s.T(), http.StatusForbidden, w.Code)
var resp map[string]interface{}
s.Require().NoError(json.Unmarshal(w.Body.Bytes(), &resp))
assert.Equal(s.T(), "Custom tools are not enabled for this account", resp["error"])
}
// ========== 获取自定义工具测试 ==========
func (s *CaptainCustomToolCRUDTestSuite) TestGet_成功获取自定义工具() {
internal/handler/api/v1/captain_custom_tool_handler.go
+26
View File
@@ -29,6 +29,9 @@ func (h *CaptainCustomToolHandler) Create(c *gin.Context) {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid account_id")
return
}
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
var req service.CreateCustomToolRequest
if err := bindNestedJSONPayload(c, "custom_tool", &req); err != nil {
@@ -54,6 +57,9 @@ func (h *CaptainCustomToolHandler) Get(c *gin.Context) {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid account_id")
return
}
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
id, err := parseUintAnyParam(c, "tool_id", "id")
if err != nil {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid id")
@@ -78,6 +84,9 @@ func (h *CaptainCustomToolHandler) Update(c *gin.Context) {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid account_id")
return
}
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
id, err := parseUintAnyParam(c, "tool_id", "id")
if err != nil {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid id")
@@ -108,6 +117,9 @@ func (h *CaptainCustomToolHandler) Delete(c *gin.Context) {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid account_id")
return
}
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
id, err := parseUintAnyParam(c, "tool_id", "id")
if err != nil {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid id")
@@ -131,6 +143,9 @@ func (h *CaptainCustomToolHandler) List(c *gin.Context) {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid account_id")
return
}
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
tools, count, err := h.svc.List(c.Request.Context(), accountID, 0, 1000)
if err != nil {
@@ -179,6 +194,9 @@ func (h *CaptainCustomToolHandler) TestTool(c *gin.Context) {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid account_id")
return
}
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
var req service.TestToolRequest
if err := bindNestedJSONPayload(c, "custom_tool", &req); err != nil {
@@ -196,6 +214,14 @@ func (h *CaptainCustomToolHandler) TestTool(c *gin.Context) {
c.JSON(http.StatusOK, result)
}
func (h *CaptainCustomToolHandler) ensureCustomToolsEnabled(c *gin.Context, accountID uint) bool {
if h.svc.CustomToolsEnabled(c.Request.Context(), accountID) {
return true
}
c.JSON(http.StatusForbidden, gin.H{"error": "Custom tools are not enabled for this account"})
return false
}
func captainCustomToolPayload(tool *model.CaptainCustomTool) gin.H {
return gin.H{
"id": tool.ID,
internal/handler/api/v1/captain_custom_tool_test_handler_test.go
+3 -2
View File
@@ -50,7 +50,7 @@ func (s *CaptainCustomToolTestHandlerTestSuite) SetupSuite() {
s.Require().NoError(err)
// 创建测试账户
account := &model.Account{Name: "CustomToolTestOrg", Locale: "en", Active: true}
account := &model.Account{Name: "CustomToolTestOrg", Locale: "en", Active: true, FeatureFlags: `{"custom_tools":true}`}
s.Require().NoError(db.Create(account).Error)
s.account = account
@@ -70,7 +70,8 @@ func (s *CaptainCustomToolTestHandlerTestSuite) SetupSuite() {
// 创建 repo + service + handler
toolRepo := repository.NewCaptainCustomToolRepo(db)
svc := service.NewCaptainCustomToolService(toolRepo)
accountRepo := repository.NewAccountRepo(db)
svc := service.NewCaptainCustomToolService(toolRepo, accountRepo)
svc.SetHTTPClient(fakeCaptainToolHTTPDoer(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusCreated, Body: http.NoBody, Header: make(http.Header)}, nil
}))
internal/handler/api/v1/captain_resource_parity_handler_test.go
+4 -3
View File
@@ -46,8 +46,8 @@ func setupCaptainResourceParityTest(t *testing.T) (*gin.Engine, *gorm.DB, *model
sqlDB.Close()
})
account := &model.Account{Name: "Captain Account", Locale: "en", Active: true}
otherAccount := &model.Account{Name: "Other Account", Locale: "en", Active: true}
account := &model.Account{Name: "Captain Account", Locale: "en", Active: true, FeatureFlags: `{"custom_tools":true}`}
otherAccount := &model.Account{Name: "Other Account", Locale: "en", Active: true, FeatureFlags: `{"custom_tools":true}`}
require.NoError(t, db.Create(account).Error)
require.NoError(t, db.Create(otherAccount).Error)
assistant := &model.CaptainAssistant{AccountID: account.ID, Name: "Fin", Description: "Support", Config: json.RawMessage(`{}`), Status: model.AssistantStatusActive}
@@ -59,7 +59,8 @@ func setupCaptainResourceParityTest(t *testing.T) (*gin.Engine, *gorm.DB, *model
scenarioHandler := NewCaptainScenarioHandler(scenarioSvc)
toolRepo := repository.NewCaptainCustomToolRepo(db)
toolSvc := service.NewCaptainCustomToolService(toolRepo)
accountRepo := repository.NewAccountRepo(db)
toolSvc := service.NewCaptainCustomToolService(toolRepo, accountRepo)
toolSvc.SetHTTPClient(fakeCaptainToolHTTPDoer(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusCreated, Body: io.NopCloser(strings.NewReader(`{"ok":true}`)), Header: make(http.Header)}, nil
}))