From 1571abd98a4057535fd31fd229a73f57d9b00771 Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Mon, 27 Jul 2026 11:23:06 +0800 Subject: [PATCH] Go: add context, part7 (#17402) Signed-off-by: Jin Hai --- internal/dao/mcp.go | 29 +++++++------ internal/handler/mcp.go | 25 +++++------ internal/service/mcp.go | 82 ++++++++++++++++++------------------ internal/service/mcp_test.go | 3 +- 4 files changed, 71 insertions(+), 68 deletions(-) diff --git a/internal/dao/mcp.go b/internal/dao/mcp.go index e95fb9b0f2..c352c75c15 100644 --- a/internal/dao/mcp.go +++ b/internal/dao/mcp.go @@ -17,6 +17,7 @@ package dao import ( + "context" "errors" "fmt" "strings" @@ -45,9 +46,9 @@ func NewMCPServerDAO() *MCPServerDAO { } // GetByID returns an MCP server by ID. -func (dao *MCPServerDAO) GetByID(id string) (*entity.MCPServer, error) { +func (dao *MCPServerDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.MCPServer, error) { var server entity.MCPServer - if err := DB.Where("id = ?", id).First(&server).Error; err != nil { + if err := db.WithContext(ctx).Where("id = ?", id).First(&server).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, nil } @@ -57,9 +58,9 @@ func (dao *MCPServerDAO) GetByID(id string) (*entity.MCPServer, error) { } // ExistsByNameAndTenant returns whether an MCP server name already exists for a tenant. -func (dao *MCPServerDAO) ExistsByNameAndTenant(name, tenantID string) (bool, error) { +func (dao *MCPServerDAO) ExistsByNameAndTenant(ctx context.Context, db *gorm.DB, name, tenantID string) (bool, error) { var count int64 - if err := DB.Model(&entity.MCPServer{}). + if err := db.WithContext(ctx).Model(&entity.MCPServer{}). Where("name = ? AND tenant_id = ?", name, tenantID). Count(&count).Error; err != nil { return false, err @@ -68,16 +69,16 @@ func (dao *MCPServerDAO) ExistsByNameAndTenant(name, tenantID string) (bool, err } // CreateMCPServer creates an MCP server. -func (dao *MCPServerDAO) CreateMCPServer(server *entity.MCPServer) error { - return DB.Create(server).Error +func (dao *MCPServerDAO) CreateMCPServer(ctx context.Context, db *gorm.DB, server *entity.MCPServer) error { + return db.WithContext(ctx).Create(server).Error } // ListMCPServers returns MCP servers for a tenant with optional filtering. -func (dao *MCPServerDAO) ListMCPServers(tenantID string, ids []string, keywords string, orderby string, desc bool) ([]*entity.MCPServer, int64, error) { +func (dao *MCPServerDAO) ListMCPServers(ctx context.Context, db *gorm.DB, tenantID string, ids []string, keywords string, orderby string, desc bool) ([]*entity.MCPServer, int64, error) { var servers []*entity.MCPServer var total int64 - query := DB.Model(&entity.MCPServer{}).Where("tenant_id = ?", tenantID) + query := db.WithContext(ctx).Model(&entity.MCPServer{}).Where("tenant_id = ?", tenantID) if len(ids) > 0 { query = query.Where("id IN ?", ids) @@ -111,17 +112,17 @@ func (dao *MCPServerDAO) ListMCPServers(tenantID string, ids []string, keywords } // GetByIDAndTenant returns an MCP server owned by a tenant. -func (dao *MCPServerDAO) GetByIDAndTenant(id, tenantID string) (*entity.MCPServer, error) { +func (dao *MCPServerDAO) GetByIDAndTenant(ctx context.Context, db *gorm.DB, id, tenantID string) (*entity.MCPServer, error) { var server entity.MCPServer - if err := DB.Where("id = ? AND tenant_id = ?", id, tenantID).First(&server).Error; err != nil { + if err := db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&server).Error; err != nil { return nil, err } return &server, nil } // DeleteMCPServer deletes an MCP server owned by a tenant. -func (dao *MCPServerDAO) DeleteMCPServer(id, tenantID string) (bool, error) { - result := DB.Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.MCPServer{}) +func (dao *MCPServerDAO) DeleteMCPServer(ctx context.Context, db *gorm.DB, id, tenantID string) (bool, error) { + result := db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.MCPServer{}) if result.Error != nil { return false, result.Error } @@ -129,8 +130,8 @@ func (dao *MCPServerDAO) DeleteMCPServer(id, tenantID string) (bool, error) { } // UpdateMCPServer updates an MCP server owned by a tenant. -func (dao *MCPServerDAO) UpdateMCPServer(id, tenantID string, updates map[string]interface{}) (bool, error) { - result := DB.Model(&entity.MCPServer{}). +func (dao *MCPServerDAO) UpdateMCPServer(ctx context.Context, db *gorm.DB, id, tenantID string, updates map[string]interface{}) (bool, error) { + result := db.WithContext(ctx).Model(&entity.MCPServer{}). Where("id = ? AND tenant_id = ?", id, tenantID). Updates(updates) if result.Error != nil { diff --git a/internal/handler/mcp.go b/internal/handler/mcp.go index 6beacb8939..3e3735259e 100644 --- a/internal/handler/mcp.go +++ b/internal/handler/mcp.go @@ -80,8 +80,9 @@ func (h *MCPHandler) CreateMCPServer(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return } + ctx := c.Request.Context() - result, code, err := h.mcpService.CreateMCPServer(user.ID, req) + result, code, err := h.mcpService.CreateMCPServer(ctx, user.ID, req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -108,13 +109,13 @@ func (h *MCPHandler) ListMCPServers(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return } - + ctx := c.Request.Context() orderby := c.DefaultQuery("orderby", "create_time") desc := strings.ToLower(c.DefaultQuery("desc", "true")) != "false" keywords := c.Query("keywords") mcpIDs := getMCPIDsFromQuery(c) - result, code, err := h.mcpService.ListMCPServers(user.ID, mcpIDs, keywords, page, pageSize, orderby, desc) + result, code, err := h.mcpService.ListMCPServers(ctx, user.ID, mcpIDs, keywords, page, pageSize, orderby, desc) if err != nil { if code == common.CodeServerError { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, code, nil, err.Error()) @@ -133,10 +134,10 @@ func (h *MCPHandler) GetMCPServer(c *gin.Context) { common.ErrorWithCode(c, errorCode, errorMessage) return } - + ctx := c.Request.Context() mcpID := c.Param("mcp_id") if c.Query("mode") == "download" { - result, code, err := h.mcpService.ExportMCPServer(user.ID, mcpID) + result, code, err := h.mcpService.ExportMCPServer(ctx, user.ID, mcpID) if err != nil { mcpDetailError(c, code, err) return @@ -145,7 +146,7 @@ func (h *MCPHandler) GetMCPServer(c *gin.Context) { return } - result, code, err := h.mcpService.GetMCPServer(user.ID, mcpID) + result, code, err := h.mcpService.GetMCPServer(ctx, user.ID, mcpID) if err != nil { mcpDetailError(c, code, err) return @@ -175,8 +176,8 @@ func (h *MCPHandler) UpdateMCPServer(c *gin.Context) { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return } - - result, code, err := h.mcpService.UpdateMCPServer(user.ID, mcpID, req) + ctx := c.Request.Context() + result, code, err := h.mcpService.UpdateMCPServer(ctx, user.ID, mcpID, req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -192,8 +193,8 @@ func (h *MCPHandler) DeleteMCPServer(c *gin.Context) { common.ErrorWithCode(c, errorCode, errorMessage) return } - - result, code, err := h.mcpService.DeleteMCPServer(user.ID, c.Param("mcp_id")) + ctx := c.Request.Context() + result, code, err := h.mcpService.DeleteMCPServer(ctx, user.ID, c.Param("mcp_id")) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -320,8 +321,8 @@ func (h *MCPHandler) ImportMCPServers(c *gin.Context) { // 10 s fallback when timeout <= 0. _ = json.Unmarshal(rawTimeout, &timeout) } - - results, err := h.mcpService.ImportServers(user.ID, servers, timeout) + ctx := c.Request.Context() + results, err := h.mcpService.ImportServers(ctx, user.ID, servers, timeout) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeBadRequest, nil, err.Error()) return diff --git a/internal/service/mcp.go b/internal/service/mcp.go index 24e7b337a2..82a30cd6a2 100644 --- a/internal/service/mcp.go +++ b/internal/service/mcp.go @@ -115,39 +115,39 @@ type ListMCPServersResponse struct { const maxMCPFetchTimeoutSec = 60 // CreateMCPServer creates an MCP server owned by a tenant. -func (s *MCPService) CreateMCPServer(tenantID string, req CreateMCPServerRequest) (*CreateMCPServerResponse, common.ErrorCode, error) { +func (s *MCPService) CreateMCPServer(ctx context.Context, tenantID string, req CreateMCPServerRequest) (*CreateMCPServerResponse, common.ErrorCode, error) { if req.Timeout < 0 || req.Timeout > maxMCPFetchTimeoutSec { - return nil, common.CodeDataError, errors.New("Invalid timeout.") + return nil, common.CodeDataError, errors.New("invalid timeout") } if !isValidMCPServerType(req.ServerType) { - return nil, common.CodeDataError, errors.New("Unsupported MCP server type.") + return nil, common.CodeDataError, errors.New("unsupported MCP server type") } if req.Name == "" || len([]byte(req.Name)) > mcpServerNameLimit { return nil, common.CodeDataError, fmt.Errorf("Invalid MCP name or length is %d which is large than 255.", len([]byte(req.Name))) } - exists, err := s.mcpServerDAO.ExistsByNameAndTenant(req.Name, tenantID) + exists, err := s.mcpServerDAO.ExistsByNameAndTenant(ctx, dao.DB, req.Name, tenantID) if err != nil { return nil, common.CodeServerError, err } if exists { - return nil, common.CodeDataError, errors.New("Duplicated MCP server name.") + return nil, common.CodeDataError, errors.New("duplicated MCP server name") } if req.URL == "" { - return nil, common.CodeDataError, errors.New("Invalid url.") + return nil, common.CodeDataError, errors.New("invalid url") } - if _, err := s.tenantDAO.GetByID(tenantID); err != nil { - return nil, common.CodeDataError, errors.New("Tenant not found.") + if _, err = s.tenantDAO.GetByID(tenantID); err != nil { + return nil, common.CodeDataError, errors.New("tenant not found") } headers := safeJSONMap(req.Headers) variables := safeJSONMap(req.Variables) delete(variables, "tools") - tools, err := fetchMCPTools(context.Background(), req.URL, req.ServerType, headers, variables, req.Timeout) + tools, err := fetchMCPTools(ctx, req.URL, req.ServerType, headers, variables, req.Timeout) if err != nil { return nil, common.CodeDataError, err } @@ -164,8 +164,8 @@ func (s *MCPService) CreateMCPServer(tenantID string, req CreateMCPServerRequest Headers: headers, } - if err := s.mcpServerDAO.CreateMCPServer(server); err != nil { - return nil, common.CodeDataError, errors.New("Failed to create MCP server.") + if err = s.mcpServerDAO.CreateMCPServer(ctx, dao.DB, server); err != nil { + return nil, common.CodeDataError, errors.New("failed to create MCP server") } return &CreateMCPServerResponse{ @@ -231,8 +231,8 @@ func jsonMapStringValues(values entity.JSONMap) map[string]string { return out } -func (s *MCPService) GetMCPServer(tenantID, mcpID string) (*entity.MCPServer, common.ErrorCode, error) { - server, err := s.mcpServerDAO.GetByIDAndTenant(mcpID, tenantID) +func (s *MCPService) GetMCPServer(ctx context.Context, tenantID, mcpID string) (*entity.MCPServer, common.ErrorCode, error) { + server, err := s.mcpServerDAO.GetByIDAndTenant(ctx, dao.DB, mcpID, tenantID) if err != nil { if isMCPServerNotFound(err) { return nil, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID) @@ -245,8 +245,8 @@ func (s *MCPService) GetMCPServer(tenantID, mcpID string) (*entity.MCPServer, co return server, common.CodeSuccess, nil } -func (s *MCPService) ExportMCPServer(userID, mcpID string) (*ExportMCPServerResponse, common.ErrorCode, error) { - server, code, err := s.GetMCPServer(userID, mcpID) +func (s *MCPService) ExportMCPServer(ctx context.Context, userID, mcpID string) (*ExportMCPServerResponse, common.ErrorCode, error) { + server, code, err := s.GetMCPServer(ctx, userID, mcpID) if err != nil { return nil, code, err } @@ -281,8 +281,8 @@ func newExportMCPServerResponse(server *entity.MCPServer) *ExportMCPServerRespon } // UpdateMCPServer updates an MCP server owned by a tenant. -func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServerRequest) (*entity.MCPServer, common.ErrorCode, error) { - server, err := s.mcpServerDAO.GetByIDAndTenant(mcpID, tenantID) +func (s *MCPService) UpdateMCPServer(ctx context.Context, tenantID, mcpID string, req UpdateMCPServerRequest) (*entity.MCPServer, common.ErrorCode, error) { + server, err := s.mcpServerDAO.GetByIDAndTenant(ctx, dao.DB, mcpID, tenantID) if err != nil { if isMCPServerNotFound(err) { return nil, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID) @@ -302,7 +302,7 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer serverTypeProvided = true } if serverTypeProvided && !isValidMCPServerType(serverType) { - return nil, common.CodeDataError, errors.New("Unsupported MCP server type.") + return nil, common.CodeDataError, errors.New("unsupported MCP server type") } serverName := server.Name @@ -324,12 +324,12 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer } else if ok { serverURL = strings.TrimSpace(value) if serverURL == "" { - return nil, common.CodeDataError, errors.New("Invalid url.") + return nil, common.CodeDataError, errors.New("invalid url") } serverURLProvided = true } if serverURL == "" { - return nil, common.CodeDataError, errors.New("Invalid url.") + return nil, common.CodeDataError, errors.New("invalid url") } headers := server.Headers @@ -355,7 +355,7 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer if err != nil { return nil, common.CodeDataError, err } - tools, err := fetchMCPTools(context.Background(), serverURL, serverType, headers, variables, timeoutSeconds) + tools, err := fetchMCPTools(ctx, serverURL, serverType, headers, variables, timeoutSeconds) if err != nil { return nil, common.CodeDataError, err } @@ -387,11 +387,11 @@ func (s *MCPService) UpdateMCPServer(tenantID, mcpID string, req UpdateMCPServer updates["description"] = description } - if _, err := s.mcpServerDAO.UpdateMCPServer(mcpID, tenantID, updates); err != nil { + if _, err = s.mcpServerDAO.UpdateMCPServer(ctx, dao.DB, mcpID, tenantID, updates); err != nil { return nil, common.CodeServerError, err } - updatedServer, err := s.mcpServerDAO.GetByIDAndTenant(mcpID, tenantID) + updatedServer, err := s.mcpServerDAO.GetByIDAndTenant(ctx, dao.DB, mcpID, tenantID) if err != nil { if isMCPServerNotFound(err) { return nil, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID) @@ -419,8 +419,8 @@ func isMCPServerNotFound(err error) bool { } // ListMCPServers lists MCP servers owned by a tenant. -func (s *MCPService) ListMCPServers(tenantID string, ids []string, keywords string, page, pageSize int, orderby string, desc bool) (*ListMCPServersResponse, common.ErrorCode, error) { - servers, total, err := s.mcpServerDAO.ListMCPServers(tenantID, ids, keywords, orderby, desc) +func (s *MCPService) ListMCPServers(ctx context.Context, tenantID string, ids []string, keywords string, page, pageSize int, orderby string, desc bool) (*ListMCPServersResponse, common.ErrorCode, error) { + servers, total, err := s.mcpServerDAO.ListMCPServers(ctx, dao.DB, tenantID, ids, keywords, orderby, desc) if err != nil { var orderbyErr *dao.InvalidMCPServerOrderByError if errors.As(err, &orderbyErr) { @@ -458,8 +458,8 @@ func (s *MCPService) ListMCPServers(tenantID string, ids []string, keywords stri } // DeleteMCPServer deletes an MCP server owned by a tenant. -func (s *MCPService) DeleteMCPServer(tenantID, mcpID string) (bool, common.ErrorCode, error) { - server, err := s.mcpServerDAO.GetByID(mcpID) +func (s *MCPService) DeleteMCPServer(ctx context.Context, tenantID, mcpID string) (bool, common.ErrorCode, error) { + server, err := s.mcpServerDAO.GetByID(ctx, dao.DB, mcpID) if err != nil { return false, common.CodeServerError, fmt.Errorf("failed to get MCP server %s: %w", mcpID, err) } @@ -467,7 +467,7 @@ func (s *MCPService) DeleteMCPServer(tenantID, mcpID string) (bool, common.Error return false, common.CodeDataError, mcpServerNotFoundError(mcpID, tenantID) } - deleted, err := s.mcpServerDAO.DeleteMCPServer(mcpID, tenantID) + deleted, err := s.mcpServerDAO.DeleteMCPServer(ctx, dao.DB, mcpID, tenantID) if err != nil { return false, common.CodeServerError, err } @@ -479,7 +479,7 @@ func (s *MCPService) DeleteMCPServer(tenantID, mcpID string) (bool, common.Error } func mcpServerNotFoundError(mcpID, tenantID string) error { - return fmt.Errorf("Cannot find MCP server %s for user %s", mcpID, tenantID) + return fmt.Errorf("cannot find MCP server %s for user %s", mcpID, tenantID) } func isValidMCPServerType(serverType string) bool { @@ -519,7 +519,7 @@ func safeJSONMap(raw json.RawMessage) entity.JSONMap { var value map[string]interface{} if err := json.Unmarshal(raw, &value); err == nil && value != nil { - return entity.JSONMap(value) + return value } var textValue string @@ -530,7 +530,7 @@ func safeJSONMap(raw json.RawMessage) entity.JSONMap { if err := json.Unmarshal([]byte(textValue), &value); err != nil || value == nil { return entity.JSONMap{} } - return entity.JSONMap(value) + return value } // ---------- import + test (this PR's additions) ---------- @@ -556,7 +556,7 @@ type ImportResult struct { } // ImportServers bulk-imports MCP servers from a {"mcpServers": {name: config}} map. -func (s *MCPService) ImportServers(tenantID string, servers map[string]map[string]interface{}, timeoutSeconds float64) ([]ImportResult, error) { +func (s *MCPService) ImportServers(ctx context.Context, tenantID string, servers map[string]map[string]interface{}, timeoutSeconds float64) ([]ImportResult, error) { if timeoutSeconds <= 0 { timeoutSeconds = defaultMCPFetchTimeoutSec } @@ -580,7 +580,7 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin } baseName := serverName - newName, err := s.nextAvailableMCPName(baseName, tenantID) + newName, err := s.nextAvailableMCPName(ctx, baseName, tenantID) if err != nil { return nil, err } @@ -627,8 +627,8 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin } } - ctx, cancel := context.WithTimeout(context.Background(), timeout) - tools, fetchErr := utility.FetchTools(ctx, utility.FetchOptions{ + mcpCtx, cancel := context.WithTimeout(context.Background(), timeout) + tools, fetchErr := utility.FetchTools(mcpCtx, utility.FetchOptions{ URL: url, ServerType: stype, Headers: headers, @@ -651,7 +651,7 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin Variables: entity.JSONMap(variables), Headers: entity.JSONMap(headerVals), } - if err := s.mcpServerDAO.CreateMCPServer(server); err != nil { + if err = s.mcpServerDAO.CreateMCPServer(ctx, dao.DB, server); err != nil { results = append(results, ImportResult{Server: serverName, Success: false, Message: "Failed to create MCP server."}) continue } @@ -665,11 +665,11 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin return results, nil } -func (s *MCPService) nextAvailableMCPName(base, tenantID string) (string, error) { +func (s *MCPService) nextAvailableMCPName(ctx context.Context, base, tenantID string) (string, error) { name := base counter := 0 for { - exists, err := s.mcpServerDAO.ExistsByNameAndTenant(name, tenantID) + exists, err := s.mcpServerDAO.ExistsByNameAndTenant(ctx, dao.DB, name, tenantID) if err != nil { return "", err } @@ -696,7 +696,7 @@ type TestServerRequest struct { // TestServer opens a live MCP session and returns the tools the server advertises. func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[string]interface{}, error) { if req == nil || req.URL == "" { - return nil, fmt.Errorf("%w: Invalid MCP url.", ErrMCPInvalidURL) + return nil, fmt.Errorf("%w: Invalid MCP url", ErrMCPInvalidURL) } if !isValidMCPServerType(req.ServerType) { return nil, ErrMCPInvalidType @@ -730,9 +730,9 @@ func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[str } } - ctx, cancel := context.WithTimeout(context.Background(), timeout) + mcpCtx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() - tools, err := utility.FetchTools(ctx, utility.FetchOptions{ + tools, err := utility.FetchTools(mcpCtx, utility.FetchOptions{ URL: req.URL, ServerType: req.ServerType, Headers: headers, diff --git a/internal/service/mcp_test.go b/internal/service/mcp_test.go index cdbece62b9..b6135d4787 100644 --- a/internal/service/mcp_test.go +++ b/internal/service/mcp_test.go @@ -67,7 +67,8 @@ func TestImportServersValidationErrors(t *testing.T) { "missing-fields": {"foo": "bar"}, "bad-type": {"url": "http://example.com", "type": "stdio"}, } - results, err := s.ImportServers("tenant-1", configs, 1) + ctx := t.Context() + results, err := s.ImportServers(ctx, "tenant-1", configs, 1) if err != nil { t.Fatalf("unexpected error: %v", err) }